@aztec/bb.js 0.0.0-devnet

package/src/barretenberg_wasm/barretenberg_wasm_base/index.ts

@@ -0,0 +1,138 @@
+import createDebug from 'debug';
+import { randomBytes } from '../../random/index.js';
+
+const debug = createDebug('bb.js:wasm');
+
+/**
+ * Base implementation of BarretenbergWasm.
+ * Contains code that is common to the "main thread" implementation and the "child thread" implementation.
+ */
+export class BarretenbergWasmBase {
+  protected memStore: { [key: string]: Uint8Array } = {};
+  protected memory!: WebAssembly.Memory;
+  protected instance!: WebAssembly.Instance;
+  protected logger: (msg: string) => void = debug;
+
+  protected getImportObj(memory: WebAssembly.Memory) {
+    /* eslint-disable camelcase */
+    const importObj = {
+      // We need to implement a part of the wasi api:
+      // https://github.com/WebAssembly/WASI/blob/main/phases/snapshot/docs.md
+      // We literally only need to support random_get, everything else is noop implementated in barretenberg.wasm.
+      wasi_snapshot_preview1: {
+        random_get: (out: any, length: number) => {
+          out = out >>> 0;
+          const randomData = randomBytes(length);
+          const mem = this.getMemory();
+          mem.set(randomData, out);
+        },
+        clock_time_get: (a1: number, a2: number, out: number) => {
+          out = out >>> 0;
+          const ts = BigInt(new Date().getTime()) * 1000000n;
+          const view = new DataView(this.getMemory().buffer);
+          view.setBigUint64(out, ts, true);
+        },
+        proc_exit: () => {
+          this.logger('PANIC: proc_exit was called.');
+          throw new Error();
+        },
+      },
+
+      // These are functions implementations for imports we've defined are needed.
+      // The native C++ build defines these in a module called "env". We must implement TypeScript versions here.
+      env: {
+        /**
+         * The 'info' call we use for logging in C++, calls this under the hood.
+         * The native code will just print to std:err (to avoid std::cout which is used for IPC).
+         * Here we just emit the log line for the client to decide what to do with.
+         */
+        logstr: (addr: number) => {
+          const str = this.stringFromAddress(addr);
+          const m = this.getMemory();
+          const str2 = `${str} (mem: ${(m.length / (1024 * 1024)).toFixed(2)}MiB)`;
+          this.logger(str2);
+          if (str2.startsWith('WARNING:')) {
+            this.logger(new Error().stack!);
+          }
+        },
+
+        get_data: (keyAddr: number, outBufAddr: number) => {
+          const key = this.stringFromAddress(keyAddr);
+          outBufAddr = outBufAddr >>> 0;
+          const data = this.memStore[key];
+          if (!data) {
+            this.logger(`get_data miss ${key}`);
+            return;
+          }
+          // this.logger(`get_data hit ${key} size: ${data.length} dest: ${outBufAddr}`);
+          // this.logger(Buffer.from(data.slice(0, 64)).toString('hex'));
+          this.writeMemory(outBufAddr, data);
+        },
+
+        set_data: (keyAddr: number, dataAddr: number, dataLength: number) => {
+          const key = this.stringFromAddress(keyAddr);
+          dataAddr = dataAddr >>> 0;
+          this.memStore[key] = this.getMemorySlice(dataAddr, dataAddr + dataLength);
+          // this.logger(`set_data: ${key} length: ${dataLength}`);
+        },
+
+        memory,
+      },
+    };
+    /* eslint-enable camelcase */
+
+    return importObj;
+  }
+
+  public exports(): any {
+    return this.instance.exports;
+  }
+
+  /**
+   * When returning values from the WASM, use >>> operator to convert signed representation to unsigned representation.
+   */
+  public call(name: string, ...args: any) {
+    if (!this.exports()[name]) {
+      throw new Error(`WASM function ${name} not found.`);
+    }
+    try {
+      return this.exports()[name](...args) >>> 0;
+    } catch (err: any) {
+      const message = `WASM function ${name} aborted, error: ${err}`;
+      this.logger(message);
+      this.logger(err.stack);
+      throw err;
+    }
+  }
+
+  public memSize() {
+    return this.getMemory().length;
+  }
+
+  /**
+   * Returns a copy of the data, not a view.
+   */
+  public getMemorySlice(start: number, end: number) {
+    return this.getMemory().subarray(start, end).slice();
+  }
+
+  public writeMemory(offset: number, arr: Uint8Array) {
+    const mem = this.getMemory();
+    mem.set(arr, offset);
+  }
+
+  // PRIVATE METHODS
+
+  private getMemory() {
+    return new Uint8Array(this.memory.buffer);
+  }
+
+  private stringFromAddress(addr: number) {
+    addr = addr >>> 0;
+    const m = this.getMemory();
+    let i = addr;
+    for (; m[i] !== 0; ++i);
+    const textDecoder = new TextDecoder('ascii');
+    return textDecoder.decode(m.slice(addr, i));
+  }
+}

package/src/barretenberg_wasm/barretenberg_wasm_main/factory/browser/index.ts

@@ -0,0 +1,10 @@
+import debug from 'debug';
+import MainWorker from './main.worker.js';
+
+export function createMainWorker() {
+  const worker = new MainWorker();
+  const debugStr = debug.disable();
+  debug.enable(debugStr);
+  worker.postMessage({ debug: debugStr });
+  return worker;
+}

package/src/barretenberg_wasm/barretenberg_wasm_main/factory/browser/main.worker.ts

@@ -0,0 +1,15 @@
+import { expose } from 'comlink';
+import { BarretenbergWasmMain } from '../../index.js';
+import debug from 'debug';
+
+self.onmessage = function (e) {
+  if (e.data.debug) {
+    debug.enable(e.data.debug);
+  }
+};
+
+expose(new BarretenbergWasmMain());
+
+self.postMessage({ ready: true });
+
+export default null as any;

package/src/barretenberg_wasm/barretenberg_wasm_main/factory/node/index.ts

@@ -0,0 +1,18 @@
+import { Worker } from 'worker_threads';
+import { dirname } from 'path';
+import { fileURLToPath } from 'url';
+
+function getCurrentDir() {
+  if (typeof __dirname !== 'undefined') {
+    return __dirname;
+  } else {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    return dirname(fileURLToPath(import.meta.url));
+  }
+}
+
+export function createMainWorker() {
+  const __dirname = getCurrentDir();
+  return new Worker(__dirname + `/main.worker.js`);
+}

package/src/barretenberg_wasm/barretenberg_wasm_main/factory/node/main.worker.ts

@@ -0,0 +1,10 @@
+import { parentPort } from 'worker_threads';
+import { expose } from 'comlink';
+import { BarretenbergWasmMain } from '../../index.js';
+import { nodeEndpoint } from '../../../helpers/node/node_endpoint.js';
+
+if (!parentPort) {
+  throw new Error('No parentPort');
+}
+
+expose(new BarretenbergWasmMain(), nodeEndpoint(parentPort));

package/src/barretenberg_wasm/barretenberg_wasm_main/heap_allocator.ts

@@ -0,0 +1,60 @@
+import { type BarretenbergWasmMain } from './index.js';
+
+/**
+ * Keeps track of heap allocations so they can be easily freed.
+ * The WASM memory layout has 1024 bytes of unused "scratch" space at the start (addresses 0-1023).
+ * We can leverage this for IO rather than making expensive bb_malloc bb_free calls.
+ * Heap allocations will be created for input/output args that don't fit into the scratch space.
+ * Input and output args can use the same scratch space as it's assume all input reads will be performed before any
+ * output writes are performed.
+ */
+export class HeapAllocator {
+  private allocs: number[] = [];
+  private inScratchRemaining = 1024;
+  private outScratchRemaining = 1024;
+
+  constructor(private wasm: BarretenbergWasmMain) {}
+
+  copyToMemory(buffers: Uint8Array[]) {
+    return buffers.map(buf => {
+      if (buf.length <= this.inScratchRemaining) {
+        const ptr = (this.inScratchRemaining -= buf.length);
+        this.wasm.writeMemory(ptr, buf);
+        return ptr;
+      } else {
+        const ptr = this.wasm.call('bbmalloc', buf.length);
+        this.wasm.writeMemory(ptr, buf);
+        this.allocs.push(ptr);
+        return ptr;
+      }
+    });
+  }
+
+  getOutputPtrs(outLens: (number | undefined)[]) {
+    return outLens.map(len => {
+      // If the obj is variable length, we need a 4 byte ptr to write the serialized data address to.
+      // WARNING: 4 only works with WASM as it has 32 bit memory.
+      const size = len || 4;
+
+      if (size <= this.outScratchRemaining) {
+        return (this.outScratchRemaining -= size);
+      } else {
+        const ptr = this.wasm.call('bbmalloc', size);
+        this.allocs.push(ptr);
+        return ptr;
+      }
+    });
+  }
+
+  addOutputPtr(ptr: number) {
+    if (ptr >= 1024) {
+      this.allocs.push(ptr);
+    }
+  }
+
+  freeAll() {
+    for (const ptr of this.allocs) {
+      this.wasm.call('bbfree', ptr);
+    }
+  }
+}

package/src/barretenberg_wasm/barretenberg_wasm_main/index.ts

@@ -0,0 +1,138 @@
+import { type Worker } from 'worker_threads';
+import createDebug from 'debug';
+import { Remote } from 'comlink';
+import { getNumCpu, getRemoteBarretenbergWasm, getSharedMemoryAvailable } from '../helpers/index.js';
+import { createThreadWorker } from '../barretenberg_wasm_thread/factory/node/index.js';
+import { type BarretenbergWasmThreadWorker } from '../barretenberg_wasm_thread/index.js';
+import { BarretenbergWasmBase } from '../barretenberg_wasm_base/index.js';
+import { HeapAllocator } from './heap_allocator.js';
+
+const debug = createDebug('bb.js:wasm');
+
+/**
+ * This is the "main thread" implementation of BarretenbergWasm.
+ * It spawns a bunch of "child thread" implementations.
+ * In a browser context, this still runs on a worker, as it will block waiting on child threads.
+ */
+export class BarretenbergWasmMain extends BarretenbergWasmBase {
+  static MAX_THREADS = 32;
+  private workers: Worker[] = [];
+  private remoteWasms: BarretenbergWasmThreadWorker[] = [];
+  private nextWorker = 0;
+  private nextThreadId = 1;
+
+  public getNumThreads() {
+    return this.workers.length + 1;
+  }
+
+  /**
+   * Init as main thread. Spawn child threads.
+   */
+  public async init(
+    module: WebAssembly.Module,
+    threads = Math.min(getNumCpu(), BarretenbergWasmMain.MAX_THREADS),
+    logger: (msg: string) => void = debug,
+    initial = 32,
+    maximum = 2 ** 16,
+  ) {
+    this.logger = logger;
+
+    const initialMb = (initial * 2 ** 16) / (1024 * 1024);
+    const maxMb = (maximum * 2 ** 16) / (1024 * 1024);
+    const shared = getSharedMemoryAvailable();
+
+    this.logger(
+      `initial mem: ${initial} pages, ${initialMb}MiB. ` +
+        `max mem: ${maximum} pages, ${maxMb}MiB. ` +
+        `threads: ${threads}, shared: ${shared}`,
+    );
+
+    this.memory = new WebAssembly.Memory({ initial, maximum, shared });
+
+    const instance = await WebAssembly.instantiate(module, this.getImportObj(this.memory));
+
+    this.instance = instance;
+
+    // Init all global/static data.
+    this.call('_initialize');
+
+    // Create worker threads. Create 1 less than requested, as main thread counts as a thread.
+    if (threads > 1) {
+      this.logger(`creating ${threads} worker threads...`);
+      this.workers = await Promise.all(Array.from({ length: threads - 1 }).map(createThreadWorker));
+      this.remoteWasms = await Promise.all(this.workers.map(getRemoteBarretenbergWasm<BarretenbergWasmThreadWorker>));
+      await Promise.all(this.remoteWasms.map(w => w.initThread(module, this.memory)));
+    }
+    this.logger('init complete.');
+  }
+
+  /**
+   * Called on main thread. Signals child threads to gracefully exit.
+   */
+  public async destroy() {
+    await Promise.all(this.workers.map(w => w.terminate()));
+  }
+
+  protected getImportObj(memory: WebAssembly.Memory) {
+    const baseImports = super.getImportObj(memory);
+
+    /* eslint-disable camelcase */
+    return {
+      ...baseImports,
+      wasi: {
+        'thread-spawn': (arg: number) => {
+          arg = arg >>> 0;
+          const id = this.nextThreadId++;
+          const worker = this.nextWorker++ % this.remoteWasms.length;
+          // this.logger(`spawning thread ${id} on worker ${worker} with arg ${arg >>> 0}`);
+          this.remoteWasms[worker].call('wasi_thread_start', id, arg).catch(this.logger);
+          // this.remoteWasms[worker].postMessage({ msg: 'thread', data: { id, arg } });
+          return id;
+        },
+      },
+      env: {
+        ...baseImports.env,
+        env_hardware_concurrency: () => {
+          // If there are no workers (we're already running as a worker, or the main thread requested no workers)
+          // then we return 1, which should cause any algos using threading to just not create a thread.
+          return this.remoteWasms.length + 1;
+        },
+      },
+    };
+    /* eslint-enable camelcase */
+  }
+
+  callWasmExport(funcName: string, inArgs: Uint8Array[], outLens: (number | undefined)[]) {
+    const alloc = new HeapAllocator(this);
+    const inPtrs = alloc.copyToMemory(inArgs);
+    const outPtrs = alloc.getOutputPtrs(outLens);
+    this.call(funcName, ...inPtrs, ...outPtrs);
+    const outArgs = this.getOutputArgs(outLens, outPtrs, alloc);
+    alloc.freeAll();
+    return outArgs;
+  }
+
+  private getOutputArgs(outLens: (number | undefined)[], outPtrs: number[], alloc: HeapAllocator) {
+    return outLens.map((len, i) => {
+      if (len) {
+        return this.getMemorySlice(outPtrs[i], outPtrs[i] + len);
+      }
+      const slice = this.getMemorySlice(outPtrs[i], outPtrs[i] + 4);
+      const ptr = new DataView(slice.buffer, slice.byteOffset, slice.byteLength).getUint32(0, true);
+
+      // Add our heap buffer to the dealloc list.
+      alloc.addOutputPtr(ptr);
+
+      // The length will be found in the first 4 bytes of the buffer, big endian. See to_heap_buffer.
+      const lslice = this.getMemorySlice(ptr, ptr + 4);
+      const length = new DataView(lslice.buffer, lslice.byteOffset, lslice.byteLength).getUint32(0, false);
+
+      return this.getMemorySlice(ptr + 4, ptr + 4 + length);
+    });
+  }
+}
+
+/**
+ * The comlink type that asyncifies the BarretenbergWasmMain api.
+ */
+export type BarretenbergWasmMainWorker = Remote<BarretenbergWasmMain>;

package/src/barretenberg_wasm/barretenberg_wasm_thread/factory/browser/index.ts

@@ -0,0 +1,10 @@
+import debug from 'debug';
+import ThreadWorker from './thread.worker.js';
+
+export function createThreadWorker() {
+  const worker = new ThreadWorker();
+  const debugStr = debug.disable();
+  debug.enable(debugStr);
+  worker.postMessage({ debug: debugStr });
+  return worker;
+}

package/src/barretenberg_wasm/barretenberg_wasm_thread/factory/browser/thread.worker.ts

@@ -0,0 +1,15 @@
+import { expose } from 'comlink';
+import { BarretenbergWasmThread } from '../../index.js';
+import debug from 'debug';
+
+self.onmessage = function (e) {
+  if (e.data.debug) {
+    debug.enable(e.data.debug);
+  }
+};
+
+expose(new BarretenbergWasmThread());
+
+self.postMessage({ ready: true });
+
+export default null as any;

package/src/barretenberg_wasm/barretenberg_wasm_thread/factory/node/index.ts

@@ -0,0 +1,18 @@
+import { Worker } from 'worker_threads';
+import { dirname } from 'path';
+import { fileURLToPath } from 'url';
+
+function getCurrentDir() {
+  if (typeof __dirname !== 'undefined') {
+    return __dirname;
+  } else {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    return dirname(fileURLToPath(import.meta.url));
+  }
+}
+
+export function createThreadWorker() {
+  const __dirname = getCurrentDir();
+  return new Worker(__dirname + `/thread.worker.js`);
+}

package/src/barretenberg_wasm/barretenberg_wasm_thread/factory/node/thread.worker.ts

@@ -0,0 +1,10 @@
+import { parentPort } from 'worker_threads';
+import { expose } from 'comlink';
+import { BarretenbergWasmThread } from '../../index.js';
+import { nodeEndpoint } from '../../../helpers/node/node_endpoint.js';
+
+if (!parentPort) {
+  throw new Error('No parentPort');
+}
+
+expose(new BarretenbergWasmThread(), nodeEndpoint(parentPort));

package/src/barretenberg_wasm/barretenberg_wasm_thread/index.ts

@@ -0,0 +1,47 @@
+import { Remote } from 'comlink';
+import { killSelf, threadLogger } from '../helpers/index.js';
+import { BarretenbergWasmBase } from '../barretenberg_wasm_base/index.js';
+
+export class BarretenbergWasmThread extends BarretenbergWasmBase {
+  /**
+   * Init as worker thread.
+   */
+  public async initThread(module: WebAssembly.Module, memory: WebAssembly.Memory) {
+    this.logger = threadLogger() || this.logger;
+    this.memory = memory;
+    this.instance = await WebAssembly.instantiate(module, this.getImportObj(this.memory));
+  }
+
+  public destroy() {
+    killSelf();
+  }
+
+  protected getImportObj(memory: WebAssembly.Memory) {
+    const baseImports = super.getImportObj(memory);
+
+    /* eslint-disable camelcase */
+    return {
+      ...baseImports,
+      wasi: {
+        'thread-spawn': () => {
+          this.logger('PANIC: threads cannot spawn threads!');
+          this.logger(new Error().stack!);
+          killSelf();
+        },
+      },
+
+      // These are functions implementations for imports we've defined are needed.
+      // The native C++ build defines these in a module called "env". We must implement TypeScript versions here.
+      env: {
+        ...baseImports.env,
+        env_hardware_concurrency: () => {
+          // We return 1, which should cause any algos using threading to just not create a thread.
+          return 1;
+        },
+      },
+    };
+    /* eslint-enable camelcase */
+  }
+}
+
+export type BarretenbergWasmThreadWorker = Remote<BarretenbergWasmThread>;

package/src/barretenberg_wasm/fetch_code/browser/index.ts

@@ -0,0 +1,12 @@
+import barretenbergModule from '../../barretenberg.wasm.gz';
+import barretenbergThreadsModule from '../../barretenberg-threads.wasm.gz';
+import pako from 'pako';
+
+// Annoyingly the wasm declares if it's memory is shared or not. So now we need two wasms if we want to be
+// able to fallback on "non shared memory" situations.
+export async function fetchCode(multithreaded: boolean) {
+  const res = await fetch(multithreaded ? barretenbergThreadsModule : barretenbergModule);
+  const compressedData = await res.arrayBuffer();
+  const decompressedData = pako.ungzip(new Uint8Array(compressedData));
+  return decompressedData.buffer;
+}

package/src/barretenberg_wasm/fetch_code/browser/wasm-module.d.ts

@@ -0,0 +1,4 @@
+declare module '*.wasm.gz' {
+  const content: string;
+  export default content;
+}

package/src/barretenberg_wasm/fetch_code/index.ts

@@ -0,0 +1 @@
+export * from './node/index.js';

package/src/barretenberg_wasm/fetch_code/node/index.ts

@@ -0,0 +1,22 @@
+import { readFile } from 'fs/promises';
+import { dirname } from 'path';
+import { fileURLToPath } from 'url';
+import pako from 'pako';
+
+function getCurrentDir() {
+  if (typeof __dirname !== 'undefined') {
+    return __dirname;
+  } else {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    return dirname(fileURLToPath(import.meta.url));
+  }
+}
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export async function fetchCode(multithreaded: boolean) {
+  const path = getCurrentDir() + '/../../barretenberg-threads.wasm.gz';
+  const compressedData = await readFile(path);
+  const decompressedData = pako.ungzip(new Uint8Array(compressedData));
+  return decompressedData.buffer;
+}

package/src/barretenberg_wasm/helpers/browser/index.ts

@@ -0,0 +1,22 @@
+import { wrap } from 'comlink';
+
+export function getSharedMemoryAvailable() {
+  const globalScope = typeof window !== 'undefined' ? window : globalThis;
+  return typeof SharedArrayBuffer !== 'undefined' && globalScope.crossOriginIsolated;
+}
+
+export function getRemoteBarretenbergWasm<T>(worker: Worker) {
+  return wrap(worker) as T;
+}
+
+export function getNumCpu() {
+  return navigator.hardwareConcurrency;
+}
+
+export function threadLogger(): ((msg: string) => void) | undefined {
+  return undefined;
+}
+
+export function killSelf() {
+  self.close();
+}

package/src/barretenberg_wasm/helpers/index.ts

@@ -0,0 +1 @@
+export * from './node/index.js';

package/src/barretenberg_wasm/helpers/node/index.ts

@@ -0,0 +1,44 @@
+import { Worker } from 'worker_threads';
+import os from 'os';
+import { wrap } from 'comlink';
+import { nodeEndpoint } from './node_endpoint.js';
+import { writeSync } from 'fs';
+
+export function getSharedMemoryAvailable() {
+  return true;
+}
+
+/**
+ * Comlink allows you to produce a Proxy to the worker, enabling you to call methods as if it were a normal class.
+ * Note we give it the type information it needs so the returned Proxy object looks like that type.
+ * Node has a different implementation, needing this nodeEndpoint wrapper, hence this function exists here.
+ */
+export function getRemoteBarretenbergWasm<T>(worker: Worker): T {
+  return wrap(nodeEndpoint(worker)) as T;
+}
+
+/**
+ * Returns number of cpus as reported by the system, unless overriden by HARDWARE_CONCURRENCY env var.
+ */
+export function getNumCpu() {
+  return +process.env.HARDWARE_CONCURRENCY! || os.cpus().length;
+}
+
+/**
+ * In node, the message passing is different to the browser. When using 'debug' in the browser, we seemingly always
+ * get our logs, but in node it looks like it's dependent on the chain of workers from child to main thread be
+ * unblocked. If one of our threads aborts, we can't see it as the parent is blocked waiting on threads to join.
+ * To work around this in node, threads will by default write directly to stdout.
+ */
+export function threadLogger(): ((msg: string) => void) | undefined {
+  return (msg: string) => {
+    writeSync(1, msg + '\n');
+  };
+}
+
+export function killSelf(): never {
+  // Extordinarily hard process termination. Due to how parent threads block on child threads etc, even process.exit
+  // doesn't seem to be able to abort the process. The following does.
+  process.kill(process.pid);
+  throw new Error();
+}

package/src/barretenberg_wasm/helpers/node/node_endpoint.ts

@@ -0,0 +1,28 @@
+import { NodeEndpoint } from 'comlink/dist/esm/node-adapter.js';
+
+export function nodeEndpoint(nep: NodeEndpoint) {
+  const listeners = new WeakMap();
+  return {
+    postMessage: nep.postMessage.bind(nep),
+    addEventListener: (_: any, eh: any) => {
+      const l = (data: any) => {
+        if ('handleEvent' in eh) {
+          eh.handleEvent({ data });
+        } else {
+          eh({ data });
+        }
+      };
+      nep.on('message', l);
+      listeners.set(eh, l);
+    },
+    removeEventListener: (_: any, eh: any) => {
+      const l = listeners.get(eh);
+      if (!l) {
+        return;
+      }
+      nep.off('message', l);
+      listeners.delete(eh);
+    },
+    start: nep.start && nep.start.bind(nep),
+  };
+}