Utility to simplify Web Worker use

useWorker.ts

type Callable = (...args: any[]) => any | Promise<any>;
type Streamable = (...args: any[]) => AsyncIterableIterator<any>;

interface Handler {
  [key: string]: any;
}

interface Callbacks {
  resolve: (value: any) => void;
  reject: (reason?: any) => void;
}

const delegatedSet = new WeakSet<any>();
const delegatedIterSet = new WeakSet<any>();

export function delegateToWorker(fn: Callable): Callable {
  delegatedSet.add(fn);
  return fn;
}

export function delegateIterToWorker(fn: Streamable): Streamable {
  delegatedIterSet.add(fn);
  return fn;
}

export function buildWorker<T extends Handler>(handler: T, worker: Worker): T {
  const callables: (Callable | undefined)[] = [];
  const streamables: (Streamable | undefined)[] = [];
  const methods = Object.keys(handler);

  methods.forEach((k, i) => {
    if (typeof handler[k] !== "function") return;
    if (delegatedSet.has(handler[k])) {
      callables[i] = handler[k].bind(handler);
    } else if (delegatedIterSet.has(handler[k])) {
      streamables[i] = handler[k].bind(handler);
    }
  });

  let k = methods.length;

  worker.onmessage = e => {
    const [callId, fnId, ...args] = e.data;
    if (callables[fnId]) {
      Promise.resolve(callables[fnId]!(...args)).then(
        result => {
          worker.postMessage([callId, 0, result], findTransferables(result));
        },
        err => {
          worker.postMessage([callId, 1, err]);
        }
      );
    } else if (streamables[fnId]) {
      const iter = streamables[fnId]!(...args);
      callables[k] = iter.next.bind(iter);
      callables[k + 1] = (value?: any) => {
        callables[k] = undefined;
        callables[k + 1] = undefined;
        iter.return?.(value);
      };
      worker.postMessage([callId, 0, k]);
      k += 2;
    }
  };

  return handler;
}

export function useWorker<T extends Handler>(handler: T, worker: Worker): T {
  const callbacks = new Map<number, Callbacks>();
  let nextCallId = 0;

  const methods = Object.keys(handler);
  methods.forEach((k, i) => {
    if (typeof handler[k] !== "function") return;
    if (delegatedSet.has(handler[k])) {
      (handler as any)[k] = (...args: any[]) => {
        const pending = new Promise((resolve, reject) => {
          worker.postMessage(
            [nextCallId, i, ...args],
            findTransferables(...args)
          );
          callbacks.set(nextCallId, { resolve, reject });
          nextCallId += 1;
        });
        return pending;
      };
    } else if (delegatedIterSet.has(handler[k])) {
      (handler as any)[k] = (...args: any[]) => {
        const pending = new Promise<number>((resolve, reject) => {
          worker.postMessage(
            [nextCallId, i, ...args],
            findTransferables(...args)
          );
          callbacks.set(nextCallId, { resolve, reject });
          nextCallId += 1;
        });

        return {
          async next(value?: any) {
            const k = await pending;
            return new Promise((resolve, reject) => {
              worker.postMessage([nextCallId, k, value]);
              callbacks.set(nextCallId, { resolve, reject });
              nextCallId += 1;
            });
          },
          async return(value?: any) {
            const k = await pending;
            return new Promise((resolve, reject) => {
              worker.postMessage([nextCallId, k + 1, value]);
              callbacks.set(nextCallId, { resolve, reject });
              nextCallId += 1;
            });
          },
          [Symbol.asyncIterator]() {
            return this;
          },
        };
      };
    }
  });

  worker.onmessage = e => {
    const [callId, status, result] = e.data;
    const callback = callbacks.get(callId);
    if (!callback) return;
    if (status === 0) {
      callback.resolve(result);
    } else if (status === 1) {
      callback.reject(result);
    }
    callbacks.delete(callId);
  };

  return handler;
}
  
function findTransferables(...values: any[]): Transferable[] {
  const transferables: Transferable[] = [];
  values.forEach(value => {
    if (value == null) return;
    if (Array.isArray(value)) {
      value.forEach(item => {
        if (item instanceof ArrayBuffer) {
          transferables.push(item);
        } else if (item.buffer && item.buffer instanceof ArrayBuffer) {
          transferables.push(item.buffer);
        }
      });
    } else if (value instanceof ArrayBuffer) {
      transferables.push(value);
    } else if (value.buffer && value.buffer instanceof ArrayBuffer) {
      transferables.push(value.buffer);
    } else if (value instanceof ImageBitmap) {
      transferables.push(value);
    }
  });
  return transferables;
}