@u1f992/pdfdiff 0.1.0 → 0.2.0

package/src/index.ts

@@ -19,16 +19,17 @@ import * as jimp from "jimp";
 import * as mupdf from "mupdf";
 import Worker from "web-worker";
 
-import {
-  createEmptyImage,
-  isValidAlignStrategy,
-  type AlignStrategy,
-} from "./image.js";
-import { withIndex } from "./iterable.js";
-import { pageToImage } from "./pdf.js";
-import { parseHex, formatHex } from "./rgba-color.js";
-import type { Pallet } from "./diff.js";
-import type { JimpInstance } from "./jimp.js";
+import { type Pallet } from "./diff.ts";
+import { isValidAlignStrategy, type AlignStrategy } from "./image.ts";
+import { withIndex } from "./iterable.ts";
+import { parseHex, formatHex } from "./rgba-color.ts";
+import type { JimpInstance } from "./jimp.ts";
+import type {
+  InitMessage,
+  PageMessage,
+  PageResultMessage,
+  ReadyMessage,
+} from "./worker.ts";
 
 export { withIndex, isValidAlignStrategy, parseHex, formatHex };
 
@@ -38,6 +39,7 @@ type Options = {
   mask: Uint8Array | undefined;
   align: AlignStrategy;
   pallet: Pallet;
+  workers: number;
 };
 
 type Result = {
@@ -59,14 +61,113 @@ export const defaultOptions: Options = {
     deletion: [0xff, 0x57, 0x24, 0xff],
     modification: [0xff, 0xc1, 0x05, 0xff],
   },
+  workers: 1,
 };
 
+function asSharedBytes(bytes: Uint8Array): Uint8Array {
+  const isNode =
+    typeof globalThis.process !== "undefined" &&
+    !!globalThis.process.versions?.node;
+  const coiOk =
+    (globalThis as { crossOriginIsolated?: boolean }).crossOriginIsolated ===
+    true;
+  if (typeof SharedArrayBuffer !== "undefined" && (isNode || coiOk)) {
+    const sab = new SharedArrayBuffer(bytes.byteLength);
+    const view = new Uint8Array(sab);
+    view.set(bytes);
+    return view;
+  }
+  return new Uint8Array(bytes);
+}
+
+class WorkerHandle {
+  worker: InstanceType<typeof Worker>;
+  private pendingResolve:
+    | ((data: ReadyMessage | PageResultMessage) => void)
+    | null = null;
+  private pendingReject: ((reason: unknown) => void) | null = null;
+
+  constructor(url: URL) {
+    this.worker = new Worker(url, { type: "module" });
+    this.worker.addEventListener(
+      "message",
+      (e: MessageEvent<ReadyMessage | PageResultMessage>) => {
+        const resolve = this.pendingResolve;
+        this.pendingResolve = null;
+        this.pendingReject = null;
+        resolve?.(e.data);
+      },
+    );
+    this.worker.addEventListener("error", (e: ErrorEvent) => {
+      const reject = this.pendingReject;
+      this.pendingResolve = null;
+      this.pendingReject = null;
+      reject?.(e.error ?? new Error(e.message));
+    });
+  }
+
+  init(msg: InitMessage): Promise<ReadyMessage> {
+    return new Promise<ReadyMessage>((resolve, reject) => {
+      this.pendingResolve = resolve as (
+        data: ReadyMessage | PageResultMessage,
+      ) => void;
+      this.pendingReject = reject;
+      this.worker.postMessage(msg);
+    });
+  }
+
+  processPage(index: number): Promise<PageResultMessage> {
+    return new Promise<PageResultMessage>((resolve, reject) => {
+      this.pendingResolve = resolve as (
+        data: ReadyMessage | PageResultMessage,
+      ) => void;
+      this.pendingReject = reject;
+      const msg: PageMessage = { type: "page", index };
+      this.worker.postMessage(msg);
+    });
+  }
+
+  terminate() {
+    this.worker.terminate();
+  }
+}
+
+function workerUrl(): URL {
+  return new URL(
+    import.meta.url.endsWith(".ts") ? "./worker.ts" : "./worker.js",
+    import.meta.url,
+  );
+}
+
+function pageResultToResult(msg: PageResultMessage): Result {
+  return {
+    a: jimp.Jimp.fromBitmap({
+      width: msg.a.width,
+      height: msg.a.height,
+      data: new Uint8Array(msg.a.data),
+    }) as JimpInstance,
+    b: jimp.Jimp.fromBitmap({
+      width: msg.b.width,
+      height: msg.b.height,
+      data: new Uint8Array(msg.b.data),
+    }) as JimpInstance,
+    diff: jimp.Jimp.fromBitmap({
+      width: msg.diff.width,
+      height: msg.diff.height,
+      data: new Uint8Array(msg.diff.data),
+    }) as JimpInstance,
+    addition: msg.addition,
+    deletion: msg.deletion,
+    modification: msg.modification,
+  };
+}
+
 export async function* visualizeDifferences(
   a: Uint8Array,
   b: Uint8Array,
   options: Partial<Omit<Options, "pallet"> & { pallet: Partial<Pallet> }>,
 ) {
-  const mergedOptions = {
+  const merged = {
     dpi: options?.dpi ?? defaultOptions.dpi,
     alpha: options?.alpha ?? defaultOptions.alpha,
     mask: options?.mask ?? defaultOptions.mask,
@@ -77,110 +178,98 @@ export async function* visualizeDifferences(
       modification:
         options?.pallet?.modification ?? defaultOptions.pallet.modification,
     },
+    workers: options?.workers ?? defaultOptions.workers,
   };
 
-  const pdfA = mupdf.PDFDocument.openDocument(a, "application/pdf");
-  const pdfB = mupdf.PDFDocument.openDocument(b, "application/pdf");
-  const pdfMask =
-    typeof mergedOptions.mask !== "undefined"
-      ? mupdf.PDFDocument.openDocument(mergedOptions.mask, "application/pdf")
+  const probe = mupdf.PDFDocument.openDocument(a, "application/pdf");
+  const probeB = mupdf.PDFDocument.openDocument(b, "application/pdf");
+  const probeMask =
+    typeof merged.mask !== "undefined"
+      ? mupdf.PDFDocument.openDocument(merged.mask, "application/pdf")
       : new mupdf.PDFDocument();
-
   const maxPages = Math.max(
-    pdfA.countPages(),
-    pdfB.countPages(),
-    pdfMask.countPages(),
+    probe.countPages(),
+    probeB.countPages(),
+    probeMask.countPages(),
   );
+  probe.destroy();
+  probeB.destroy();
+  probeMask.destroy();
 
-  async function processPage(pageIndex: number) {
-    const [pageA, pageB, pageMask] = await Promise.all([
-      pageIndex < pdfA.countPages()
-        ? pageToImage(
-            pdfA.loadPage(pageIndex),
-            mergedOptions.dpi,
-            mergedOptions.alpha,
-          )
-        : createEmptyImage(1, 1),
-      pageIndex < pdfB.countPages()
-        ? pageToImage(
-            pdfB.loadPage(pageIndex),
-            mergedOptions.dpi,
-            mergedOptions.alpha,
-          )
-        : createEmptyImage(1, 1),
-      pageIndex < pdfMask.countPages()
-        ? pageToImage(
-            pdfMask.loadPage(pageIndex),
-            mergedOptions.dpi,
-            mergedOptions.alpha,
-          )
-        : createEmptyImage(1, 1),
-    ]);
-
-    // NOTE: getBufferはcopyなので、Workerに移譲した後もa, bを使用して問題ない
-    // https://github.com/jimp-dev/jimp/blob/b6b0e418a5f1259211a133b20cddb4f4e5c25679/packages/core/src/index.ts#L444
-    const [bufA, bufB, bufMask] = await Promise.all([
-      pageA
-        .getBuffer(jimp.JimpMime.png)
-        .then((buf) => new Uint8Array(buf).buffer),
-      pageB
-        .getBuffer(jimp.JimpMime.png)
-        .then((buf) => new Uint8Array(buf).buffer),
-      pageMask
-        .getBuffer(jimp.JimpMime.png)
-        .then((buf) => new Uint8Array(buf).buffer),
-    ]);
-
-    const { bufDiff, addition, deletion, modification } = (await new Promise(
-      (resolve, reject) => {
-        const url = new URL("./worker.js", import.meta.url);
-        const worker = new Worker(url, { type: "module" });
-        worker.addEventListener("message", (e) => {
-          resolve(e.data);
-          worker.terminate();
-        });
-        worker.addEventListener("error", (e) => {
-          reject(e);
-          worker.terminate();
-        });
-        worker.postMessage(
-          {
-            bufA,
-            bufB,
-            bufMask,
-            pallet: mergedOptions.pallet,
-            align: mergedOptions.align,
-          },
-          [bufA, bufB, bufMask],
-        );
-      },
-    )) as {
-      bufDiff: ArrayBuffer;
-      addition: [number, number][];
-      deletion: [number, number][];
-      modification: [number, number][];
-    };
-    const diff = await jimp.Jimp.fromBuffer(bufDiff);
-    return { a: pageA, b: pageB, diff, addition, deletion, modification };
+  if (maxPages === 0) return;
+
+  const aBytes = asSharedBytes(a);
+  const bBytes = asSharedBytes(b);
+  const maskBytes =
+    typeof merged.mask !== "undefined" ? asSharedBytes(merged.mask) : null;
+
+  const initMsg: InitMessage = {
+    type: "init",
+    aBytes,
+    bBytes,
+    maskBytes,
+    dpi: merged.dpi,
+    alpha: merged.alpha,
+    pallet: merged.pallet,
+    align: merged.align,
+  };
+
+  const N = Math.max(1, Math.min(merged.workers, maxPages));
+  const url = workerUrl();
+  const worker0 = new WorkerHandle(url);
+  await worker0.init(initMsg);
+
+  const buffered = new Map<number, Result>();
+  let nextToAssign = 0;
+
+  const workers: WorkerHandle[] = [worker0];
+  for (let i = 1; i < N; i++) {
+    const w = new WorkerHandle(url);
+    await w.init(initMsg);
+    workers.push(w);
   }
 
-  // ページ処理を並列発行し、順序を保証して出力
-  const concurrency = navigator.hardwareConcurrency;
-  const pending = /** @type {Promise<VisualizeDifferencesResult>[]} */ [];
-  let nextPageToProcess = 0;
-  let nextPageToYield = 0;
-
-  while (nextPageToYield < maxPages) {
-    // プールに空きがあれば新しいページ処理を追加
-    while (nextPageToProcess < maxPages && pending.length < concurrency) {
-      pending.push(processPage(nextPageToProcess));
-      nextPageToProcess++;
+  const resolvers = new Map<number, (r: Result) => void>();
+  let workerError: unknown = null;
+
+  const loops = workers.map(async (w) => {
+    while (nextToAssign < maxPages && workerError === null) {
+      const idx = nextToAssign++;
+      try {
+        const msg = await w.processPage(idx);
+        const result = pageResultToResult(msg);
+        const resolve = resolvers.get(idx);
+        if (resolve) {
+          resolvers.delete(idx);
+          resolve(result);
+        } else {
+          buffered.set(idx, result);
+        }
+      } catch (e) {
+        workerError = e;
+        for (const [, resolve] of resolvers) resolve(null as never);
+        resolvers.clear();
+        return;
+      }
     }
+  });
 
-    // 次に出力すべきページのPromiseを待つ
-    const result = await pending[0];
-    pending.shift();
-    yield result as Result;
-    nextPageToYield++;
+  try {
+    for (let i = 0; i < maxPages; i++) {
+      if (workerError !== null) throw workerError;
+      let r: Result;
+      const buf = buffered.get(i);
+      if (buf !== undefined) {
+        buffered.delete(i);
+        r = buf;
+      } else {
+        r = await new Promise<Result>((resolve) => resolvers.set(i, resolve));
+        if (workerError !== null) throw workerError;
+      }
+      yield r;
+    }
+    await Promise.all(loops);
+  } finally {
+    for (const w of workers) w.terminate();
   }
 }

package/src/jimp.ts

@@ -4,6 +4,7 @@ export type JimpInstance = Pick<
   jimp.JimpInstance,
   | "width"
   | "height"
+  | "bitmap"
   | "getPixelColor"
   | "setPixelColor"
   | "resize"

package/src/pdf.ts

@@ -18,12 +18,42 @@
 import * as jimp from "jimp";
 import * as mupdf from "mupdf";
 
+import type { JimpInstance } from "./jimp.ts";
+
 export function* loadPages(pdf: mupdf.Document) {
   for (let i = 0; i < pdf.countPages(); i++) {
     yield pdf.loadPage(i);
   }
 }
 
+function pixmapToRGBA(pixmap: mupdf.Pixmap): Uint8Array {
+  const width = pixmap.getWidth();
+  const height = pixmap.getHeight();
+  const stride = pixmap.getStride();
+  const hasAlpha = pixmap.getAlpha() !== 0;
+  const samples = pixmap.getPixels();
+
+  if (hasAlpha && stride === width * 4) {
+    return new Uint8Array(samples);
+  }
+
+  const out = new Uint8Array(width * height * 4);
+  const srcBpp = pixmap.getNumberOfComponents() + (hasAlpha ? 1 : 0);
+  for (let y = 0; y < height; y++) {
+    const srcRow = y * stride;
+    const dstRow = y * width * 4;
+    for (let x = 0; x < width; x++) {
+      const s = srcRow + x * srcBpp;
+      const d = dstRow + x * 4;
+      out[d] = samples[s]!;
+      out[d + 1] = samples[s + 1]!;
+      out[d + 2] = samples[s + 2]!;
+      out[d + 3] = hasAlpha ? samples[s + 3]! : 255;
+    }
+  }
+  return out;
+}
+
 export async function pageToImage(
   page: mupdf.Page,
   dpi: number,
@@ -35,8 +65,10 @@ export async function pageToImage(
     mupdf.ColorSpace.DeviceRGB,
     alpha,
   );
-  const ret = await jimp.Jimp.fromBuffer(new Uint8Array(pixmap.asPNG()).buffer);
+  const width = pixmap.getWidth();
+  const height = pixmap.getHeight();
+  const data = pixmapToRGBA(pixmap);
   pixmap.destroy();
   page.destroy();
-  return ret;
+  return jimp.Jimp.fromBitmap({ width, height, data }) as JimpInstance;
 }

package/src/worker.ts

@@ -15,48 +15,145 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-import * as jimp from "jimp";
+import * as mupdf from "mupdf";
 
-import { drawDifference, type Pallet } from "./diff.js";
-import { composeLayers, type AlignStrategy } from "./image.js";
-import type { JimpInstance } from "./jimp.js";
+import { drawDifference, type Pallet } from "./diff.ts";
+import {
+  composeLayers,
+  createEmptyImage,
+  type AlignStrategy,
+} from "./image.ts";
+import type { JimpInstance } from "./jimp.ts";
+import { pageToImage } from "./pdf.ts";
+
+export type InitMessage = {
+  type: "init";
+  aBytes: Uint8Array;
+  bBytes: Uint8Array;
+  maskBytes: Uint8Array | null;
+  dpi: number;
+  alpha: boolean;
+  pallet: Pallet;
+  align: AlignStrategy;
+};
+
+export type PageMessage = {
+  type: "page";
+  index: number;
+};
+
+export type ReadyMessage = {
+  type: "ready";
+};
+
+export type PageResultMessage = {
+  type: "pageResult";
+  index: number;
+  a: { width: number; height: number; data: ArrayBuffer };
+  b: { width: number; height: number; data: ArrayBuffer };
+  diff: { width: number; height: number; data: ArrayBuffer };
+  addition: [number, number][];
+  deletion: [number, number][];
+  modification: [number, number][];
+};
+
+let pdfA: mupdf.Document;
+let pdfB: mupdf.Document;
+let pdfMask: mupdf.Document;
+let opts: {
+  dpi: number;
+  alpha: boolean;
+  pallet: Pallet;
+  align: AlignStrategy;
+};
+
+function toTransferable(
+  src: Buffer | Uint8Array | Uint8ClampedArray | number[],
+): ArrayBuffer {
+  const view =
+    src instanceof Uint8Array || src instanceof Uint8ClampedArray
+      ? src
+      : Uint8Array.from(src as ArrayLike<number>);
+  const out = new ArrayBuffer(view.byteLength);
+  new Uint8Array(out).set(view);
+  return out;
+}
+
+async function processPage(index: number): Promise<PageResultMessage> {
+  const [pageA, pageB, pageMask] = (await Promise.all([
+    index < pdfA.countPages()
+      ? pageToImage(pdfA.loadPage(index), opts.dpi, opts.alpha)
+      : createEmptyImage(1, 1),
+    index < pdfB.countPages()
+      ? pageToImage(pdfB.loadPage(index), opts.dpi, opts.alpha)
+      : createEmptyImage(1, 1),
+    index < pdfMask.countPages()
+      ? pageToImage(pdfMask.loadPage(index), opts.dpi, opts.alpha)
+      : createEmptyImage(1, 1),
+  ])) as [JimpInstance, JimpInstance, JimpInstance];
+
+  const {
+    diff: diffLayer,
+    addition,
+    deletion,
+    modification,
+  } = drawDifference(pageA, pageB, pageMask, opts.pallet, opts.align);
+  const diff = composeLayers(pageA.width, pageA.height, [
+    [pageA, 0.2],
+    [pageB, 0.2],
+    [diffLayer, 1],
+  ]);
+
+  return {
+    type: "pageResult",
+    index,
+    a: {
+      width: pageA.width,
+      height: pageA.height,
+      data: toTransferable(pageA.bitmap.data),
+    },
+    b: {
+      width: pageB.width,
+      height: pageB.height,
+      data: toTransferable(pageB.bitmap.data),
+    },
+    diff: {
+      width: diff.width,
+      height: diff.height,
+      data: toTransferable(diff.bitmap.data),
+    },
+    addition,
+    deletion,
+    modification,
+  };
+}
 
 self.addEventListener(
   "message",
-  async (
-    e: MessageEvent<{
-      bufA: ArrayBuffer;
-      bufB: ArrayBuffer;
-      bufMask: ArrayBuffer;
-      pallet: Pallet;
-      align: AlignStrategy;
-    }>,
-  ) => {
-    const { bufA, bufB, bufMask, pallet, align } = e.data;
-    const a = (await jimp.Jimp.fromBuffer(bufA)) as JimpInstance;
-    const b = (await jimp.Jimp.fromBuffer(bufB)) as JimpInstance;
-    const mask = (await jimp.Jimp.fromBuffer(bufMask)) as JimpInstance;
-    const {
-      diff: diffLayer,
-      addition,
-      deletion,
-      modification,
-    } = drawDifference(a, b, mask, pallet, align);
-    const diff = composeLayers(a.width, a.height, [
-      [a, 0.2],
-      [b, 0.2],
-      [diffLayer, 1],
-    ]);
-    const bufDiff = new Uint8Array(await diff.getBuffer(jimp.JimpMime.png))
-      .buffer;
-    self.postMessage(
-      {
-        bufDiff,
-        addition,
-        deletion,
-        modification,
-      },
-      [bufDiff],
-    );
+  async (e: MessageEvent<InitMessage | PageMessage>) => {
+    const msg = e.data;
+    if (msg.type === "init") {
+      pdfA = mupdf.PDFDocument.openDocument(msg.aBytes, "application/pdf");
+      pdfB = mupdf.PDFDocument.openDocument(msg.bBytes, "application/pdf");
+      pdfMask = msg.maskBytes
+        ? mupdf.PDFDocument.openDocument(msg.maskBytes, "application/pdf")
+        : new mupdf.PDFDocument();
+      opts = {
+        dpi: msg.dpi,
+        alpha: msg.alpha,
+        pallet: msg.pallet,
+        align: msg.align,
+      };
+      if (pdfA.countPages() > 0) pdfA.loadPage(0).destroy();
+      const ready: ReadyMessage = { type: "ready" };
+      self.postMessage(ready);
+    } else if (msg.type === "page") {
+      const result = await processPage(msg.index);
+      self.postMessage(result, [
+        result.a.data,
+        result.b.data,
+        result.diff.data,
+      ]);
+    }
   },
 );