@aztec/bb.js 0.0.1-alpha.7 → 0.0.1-commit.5daedc8

package/src/barretenberg_wasm/barretenberg_wasm_main/index.ts

@@ -0,0 +1,251 @@
+import { type Worker } from 'worker_threads';
+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';
+
+/**
+ * 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;
+  private useCustomLogger = false;
+
+  // Pre-allocated scratch buffers for msgpack I/O to avoid malloc/free overhead
+  private msgpackInputScratch: number = 0; // 8MB input buffer
+  private msgpackOutputScratch: number = 0; // 8MB output buffer
+  private readonly MSGPACK_SCRATCH_SIZE = 1024 * 1024 * 8; // 8MB
+
+  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,
+    initial = 33,
+    maximum = this.getDefaultMaximumMemoryPages(),
+  ) {
+    // Track whether a custom logger was provided so workers know whether to postMessage logs
+    this.useCustomLogger = logger !== undefined;
+    this.logger = logger ?? (() => {});
+
+    const initialMb = (initial * 2 ** 16) / (1024 * 1024);
+    const maxMb = (maximum * 2 ** 16) / (1024 * 1024);
+    const shared = getSharedMemoryAvailable();
+
+    this.logger(
+      `Initializing bb wasm: initial memory ${initial} pages ${initialMb}MiB; ` +
+        `max memory: ${maximum} pages, ${maxMb}MiB; ` +
+        `threads: ${threads}; shared memory: ${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');
+
+    // Allocate dedicated msgpack scratch buffers (never freed, reused for all msgpack calls)
+    this.msgpackInputScratch = this.call('bbmalloc', this.MSGPACK_SCRATCH_SIZE);
+    this.msgpackOutputScratch = this.call('bbmalloc', this.MSGPACK_SCRATCH_SIZE);
+    this.logger(
+      `Allocated msgpack scratch buffers: ` +
+        `input @ ${this.msgpackInputScratch}, output @ ${this.msgpackOutputScratch} (${this.MSGPACK_SCRATCH_SIZE} bytes each)`,
+    );
+
+    // 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));
+
+      // Set up log message forwarding from workers to our logger (only if custom logger provided)
+      if (this.useCustomLogger) {
+        this.workers.forEach(worker => this.setupWorkerLogForwarding(worker));
+      }
+
+      this.remoteWasms = await Promise.all(this.workers.map(getRemoteBarretenbergWasm<BarretenbergWasmThreadWorker>));
+      await Promise.all(this.remoteWasms.map(w => w.initThread(module, this.memory, this.useCustomLogger)));
+    }
+  }
+
+  private getDefaultMaximumMemoryPages(): number {
+    // iOS browser is very aggressive with memory. Check if running in browser and on iOS
+    // We at any rate expect the mobile iOS browser to kill us >=1GB, so we don't set a maximum higher than that.
+    if (typeof window !== 'undefined' && /iPad|iPhone/.test(navigator.userAgent)) {
+      return 2 ** 14;
+    }
+    return 2 ** 16;
+  }
+
+  /**
+   * Set up forwarding of log messages from worker threads to our logger.
+   * Workers post messages with { type: 'log', msg: string } which we intercept here.
+   */
+  private setupWorkerLogForwarding(worker: Worker) {
+    const handler = (data: unknown) => {
+      if (data && typeof data === 'object' && 'type' in data && data.type === 'log' && 'msg' in data) {
+        this.logger(data.msg as string);
+      }
+    };
+
+    // Node Workers use 'on' method, browser Workers use 'addEventListener'
+    // The 'worker' variable is typed as Node's Worker, but at runtime in browser
+    // it will be a browser Worker (due to browser_postprocess.sh import rewriting)
+    if ('on' in worker && typeof worker.on === 'function') {
+      // Node.js worker_threads Worker
+      worker.on('message', handler);
+    } else if ('addEventListener' in worker) {
+      // Browser Web Worker
+      (worker as unknown as globalThis.Worker).addEventListener('message', (event: MessageEvent) => {
+        handler(event.data);
+      });
+    }
+  }
+
+  /**
+   * 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 | number)[], outLens: (number | undefined)[]) {
+    const alloc = new HeapAllocator(this);
+    const inPtrs = alloc.getInputs(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);
+    });
+  }
+
+  cbindCall(cbind: string, inputBuffer: Uint8Array): any {
+    const needsCustomInputBuffer = inputBuffer.length > this.MSGPACK_SCRATCH_SIZE;
+    let inputPtr: number;
+
+    if (needsCustomInputBuffer) {
+      // Allocate temporary buffer for oversized input
+      inputPtr = this.call('bbmalloc', inputBuffer.length);
+    } else {
+      // Use pre-allocated scratch buffer
+      inputPtr = this.msgpackInputScratch;
+    }
+
+    // Write input to buffer
+    this.writeMemory(inputPtr, inputBuffer);
+
+    // Setup output scratch buffer with IN-OUT parameter pattern:
+    // Reserve 8 bytes for metadata (pointer + size), rest is scratch data space
+    const METADATA_SIZE = 8;
+    const outputPtrLocation = this.msgpackOutputScratch;
+    const outputSizeLocation = this.msgpackOutputScratch + 4;
+    const scratchDataPtr = this.msgpackOutputScratch + METADATA_SIZE;
+    const scratchDataSize = this.MSGPACK_SCRATCH_SIZE - METADATA_SIZE;
+
+    // Get memory and create DataView for writing IN values
+    let mem = this.getMemory();
+    let view = new DataView(mem.buffer);
+
+    // Write IN values: provide scratch buffer pointer and size to C++
+    view.setUint32(outputPtrLocation, scratchDataPtr, true);
+    view.setUint32(outputSizeLocation, scratchDataSize, true);
+
+    // Call WASM
+    this.call(cbind, inputPtr, inputBuffer.length, outputPtrLocation, outputSizeLocation);
+
+    // Free custom input buffer if allocated
+    if (needsCustomInputBuffer) {
+      this.call('bbfree', inputPtr);
+    }
+
+    // Re-fetch memory after WASM call, as the buffer may have been detached if memory grew
+    mem = this.getMemory();
+    view = new DataView(mem.buffer);
+
+    // Read OUT values: C++ returns actual buffer pointer and size
+    const outputDataPtr = view.getUint32(outputPtrLocation, true);
+    const outputSize = view.getUint32(outputSizeLocation, true);
+
+    // Check if C++ used scratch (pointer unchanged) or allocated (pointer changed)
+    const usedScratch = outputDataPtr === scratchDataPtr;
+
+    // Copy output data from WASM memory
+    const encodedResult = this.getMemorySlice(outputDataPtr, outputDataPtr + outputSize);
+
+    // Only free if C++ allocated beyond scratch
+    if (!usedScratch) {
+      this.call('bbfree', outputDataPtr);
+    }
+
+    return encodedResult;
+  }
+}
+
+/**
+ * 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,9 @@
+import { readinessListener } from '../../../helpers/browser/index.js';
+
+export async function createThreadWorker() {
+  // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+  // @ts-ignore
+  const worker = new Worker(new URL('./thread.worker.js', import.meta.url), { type: 'module' });
+  await new Promise<void>(resolve => readinessListener(worker, resolve));
+  return worker;
+}

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

@@ -0,0 +1,6 @@
+import { expose } from 'comlink';
+import { BarretenbergWasmThread } from '../../index.js';
+import { Ready } from '../../../helpers/browser/index.js';
+
+expose(new BarretenbergWasmThread());
+postMessage(Ready);

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

@@ -0,0 +1,19 @@
+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();
+  const worker = new Worker(__dirname + `/thread.worker.js`);
+  return Promise.resolve(worker);
+}

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

@@ -0,0 +1,12 @@
+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');
+}
+
+const endpoint = nodeEndpoint(parentPort);
+
+expose(new BarretenbergWasmThread(), endpoint);

package/src/barretenberg_wasm/barretenberg_wasm_thread/index.ts

@@ -0,0 +1,48 @@
+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.
+   * @param useCustomLogger - If true, logs will be posted back to main thread for custom logger routing
+   */
+  public async initThread(module: WebAssembly.Module, memory: WebAssembly.Memory, useCustomLogger = false) {
+    this.logger = threadLogger(useCustomLogger) || 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/barretenberg-threads.ts

@@ -0,0 +1,3 @@
+import barretenbergThreadsModule from '../../barretenberg-threads.wasm.gz';
+
+export default barretenbergThreadsModule;

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

@@ -0,0 +1,3 @@
+import barretenbergModule from '../../barretenberg.wasm.gz';
+
+export default barretenbergModule;

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

@@ -0,0 +1,34 @@
+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, wasmPath?: string) {
+  let url: string;
+  if (wasmPath) {
+    const suffix = multithreaded ? '-threads' : '';
+    const filePath = wasmPath.split('/').slice(0, -1).join('/');
+    const fileNameWithExtensions = wasmPath.split('/').pop();
+    const [fileName, ...extensions] = fileNameWithExtensions!.split('.');
+    url = `${filePath}/${fileName}${suffix}.${extensions.join('.')}`;
+  } else {
+    url = multithreaded
+      ? (await import('./barretenberg-threads.js')).default
+      : (await import('./barretenberg.js')).default;
+  }
+  const res = await fetch(url);
+  // Default bb wasm is compressed, but user could point it to a non-compressed version
+  const maybeCompressedData = await res.arrayBuffer();
+  const buffer = new Uint8Array(maybeCompressedData);
+  const isGzip =
+    // Check magic number
+    buffer[0] === 0x1f &&
+    buffer[1] === 0x8b &&
+    // Check compression method:
+    buffer[2] === 0x08;
+  if (isGzip) {
+    const decompressedData = pako.ungzip(buffer);
+    return decompressedData.buffer as unknown as Uint8Array<ArrayBuffer>;
+  } else {
+    return buffer;
+  }
+}

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,34 @@
+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, wasmPath?: string) {
+  const path = wasmPath ?? getCurrentDir() + '/../../barretenberg-threads.wasm.gz';
+  // Default bb wasm is compressed, but user could point it to a non-compressed version
+  const maybeCompressedData = await readFile(path);
+  const buffer = new Uint8Array(maybeCompressedData);
+  const isGzip =
+    // Check magic number
+    buffer[0] === 0x1f &&
+    buffer[1] === 0x8b &&
+    // Check compression method:
+    buffer[2] === 0x08;
+  if (isGzip) {
+    const decompressedData = pako.ungzip(buffer);
+    return decompressedData.buffer as unknown as Uint8Array<ArrayBuffer>;
+  } else {
+    return buffer;
+  }
+}

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

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

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

@@ -0,0 +1,54 @@
+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<T>(worker);
+}
+
+export function getNumCpu() {
+  return navigator.hardwareConcurrency;
+}
+
+export function threadLogger(useCustomLogger: boolean): ((msg: string) => void) | undefined {
+  if (useCustomLogger) {
+    // Post log messages back to main thread for routing through user-provided logger
+    return (msg: string) => {
+      postMessage({ type: 'log', msg });
+    };
+  }
+  // Use console.log directly when no custom logger is provided
+  return console.log;
+}
+
+export function killSelf() {
+  self.close();
+}
+
+export function getAvailableThreads(logger: (msg: string) => void): number {
+  if (typeof navigator !== 'undefined' && navigator.hardwareConcurrency) {
+    return navigator.hardwareConcurrency;
+  } else {
+    logger(`Could not detect environment to query number of threads. Falling back to one thread.`);
+    return 1;
+  }
+}
+
+// Solution to async initialization of workers, taken from
+// https://github.com/GoogleChromeLabs/comlink/issues/635#issuecomment-1598913044
+
+/** The message expected by the `readinessListener`. */
+export const Ready = { ready: true };
+
+/** Listen for the readiness message from the Worker and call the `callback` once. */
+export function readinessListener(worker: Worker, callback: () => void) {
+  worker.addEventListener('message', function ready(event: MessageEvent<typeof Ready>) {
+    if (!!event.data && event.data.ready === true) {
+      worker.removeEventListener('message', ready);
+      callback();
+    }
+  });
+}

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,62 @@
+import { Worker, parentPort } 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) {
+  return wrap<T>(nodeEndpoint(worker));
+}
+
+/**
+ * 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;
+}
+
+/**
+ * Returns a logger function for worker threads.
+ * When a custom logger is provided, posts messages back to the main thread.
+ * Otherwise, writes directly to stdout.
+ */
+export function threadLogger(useCustomLogger: boolean): ((msg: string) => void) | undefined {
+  if (useCustomLogger) {
+    return (msg: string) => {
+      if (parentPort) {
+        parentPort.postMessage({ type: 'log', msg });
+      }
+    };
+  }
+  // Write directly to stdout when no custom logger is provided
+  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();
+}
+
+export function getAvailableThreads(logger: (msg: string) => void): number {
+  try {
+    return os.cpus().length;
+  } catch (e: any) {
+    logger(
+      `Could not detect environment to query number of threads. Falling back to one thread. Error: ${e.message ?? e}`,
+    );
+    return 1;
+  }
+}

package/src/barretenberg_wasm/index.test.ts

@@ -0,0 +1,46 @@
+import { createMainWorker } from '../barretenberg_wasm/barretenberg_wasm_main/factory/node/index.js';
+import { BarretenbergWasmMainWorker } from '../barretenberg_wasm/barretenberg_wasm_main/index.js';
+import { getRemoteBarretenbergWasm } from '../barretenberg_wasm/helpers/index.js';
+import { fetchModuleAndThreads } from '../barretenberg_wasm/index.js';
+import { Worker } from 'worker_threads';
+
+describe('barretenberg wasm', () => {
+  let wasm: BarretenbergWasmMainWorker;
+  let worker: Worker;
+
+  beforeAll(async () => {
+    worker = await createMainWorker();
+    wasm = getRemoteBarretenbergWasm<BarretenbergWasmMainWorker>(worker);
+    const { module, threads } = await fetchModuleAndThreads(2);
+    await wasm.init(module, threads);
+  }, 20000);
+
+  afterAll(async () => {
+    await wasm.destroy();
+    await worker.terminate();
+  });
+
+  it('should new malloc, transfer and slice mem', async () => {
+    const length = 1024;
+    const ptr = await wasm.call('bbmalloc', length);
+    const buf = Buffer.alloc(length, 128);
+    await wasm.writeMemory(ptr, Uint8Array.from(buf));
+    const result = Buffer.from(await wasm.getMemorySlice(ptr, ptr + length));
+    await wasm.call('bbfree', ptr);
+    expect(result).toStrictEqual(buf);
+  });
+
+  it('test abort', async () => {
+    await expect(() => wasm.call('test_abort')).rejects.toThrow();
+  });
+
+  it('should new malloc, transfer and slice mem', async () => {
+    const length = 1024;
+    const ptr = await wasm.call('bbmalloc', length);
+    const buf = Buffer.alloc(length, 128);
+    await wasm.writeMemory(ptr, Uint8Array.from(buf));
+    const result = Buffer.from(await wasm.getMemorySlice(ptr, ptr + length));
+    await wasm.call('bbfree', ptr);
+    expect(result).toStrictEqual(buf);
+  });
+});

package/src/barretenberg_wasm/index.ts

@@ -1 +1,21 @@
-export * from './barretenberg_wasm.js';
+import { getSharedMemoryAvailable, getAvailableThreads } from './helpers/node/index.js';
+import { fetchCode } from './fetch_code/index.js';
+
+export async function fetchModuleAndThreads(
+  desiredThreads = 32,
+  wasmPath?: string,
+  logger: (msg: string) => void = () => {},
+) {
+  const shared = getSharedMemoryAvailable();
+
+  const availableThreads = shared ? await getAvailableThreads(logger) : 1;
+  // We limit the number of threads to 32 as we do not benefit from greater numbers.
+  const limitedThreads = Math.min(desiredThreads, availableThreads, 32);
+
+  logger(`Fetching bb wasm from ${wasmPath ?? 'default location'}`);
+  const code = await fetchCode(shared, wasmPath);
+  logger(`Compiling bb wasm of ${code.byteLength} bytes`);
+  const module = await WebAssembly.compile(code);
+  logger('Compilation of bb wasm complete');
+  return { module, threads: limitedThreads };
+}