@aztec/foundation 0.1.0-alpha10

package/src/wasm/wasm/empty_wasi_sdk.ts

@@ -0,0 +1,166 @@
+import { createDebugLogger } from '../../log/index.js';
+
+/**
+ * Dummy implementation of a necessary part of the wasi api:
+ * https://github.com/WebAssembly/WASI/blob/main/phases/snapshot/docs.md
+ * We don't use these functions, but the environment expects them.
+ * TODO find a way to update off of wasi 12.
+ */
+/* eslint-disable camelcase */
+export const getEmptyWasiSdk = (debug = createDebugLogger('wasm:empty_wasi_sdk')) => ({
+  /**
+   * Retrieves the current time from the system clock.
+   * This function is a dummy implementation of the WASI API's `clock_time_get` method,
+   * which is expected by the environment but not used in this context.
+   *
+   * No input parameters or return values are required, as the purpose of this function
+   * is solely to satisfy the environment expectations and provide debugging information.
+   */
+  clock_time_get() {
+    debug('clock_time_get');
+  },
+  /**
+   * Dummy implementation of WASI's environ_get function.
+   * This function is used to obtain a snapshot of the current environment variables.
+   * In this dummy implementation, no actual actions are performed, but the debug logger logs 'environ_get' when called.
+   * Environment variables are not used in this context, so the real implementation is not required.
+   *
+   * @see https://github.com/WebAssembly/WASI/blob/main/phases/snapshot/docs.md#environ_get
+   */
+  environ_get() {
+    debug('environ_get');
+  },
+  /**
+   * Retrieves the environment variable sizes from the WebAssembly environment.
+   * This function is part of the WASI API and provides a dummy implementation to fulfill the expected APIs.
+   * It does not have any actual functionality, but serves as a placeholder in the environment.
+   */
+  environ_sizes_get() {
+    debug('environ_sizes_get');
+  },
+  /**
+   * Closes a file descriptor, releasing any resources associated with it.
+   * This function does not perform any actual closing operation, but exists to
+   * satisfy the requirements of the WebAssembly System Interface (WASI) API,
+   * which expects certain functions to be present for compatibility purposes.
+   *
+   * @see https://github.com/WebAssembly/WASI/blob/main/phases/snapshot/docs.md
+   */
+  fd_close() {
+    debug('fd_close');
+  },
+  /**
+   * A dummy implementation of the 'fd_read' function from the WASI API.
+   * This function is required by the environment, but not used in this context.
+   * It would normally read data from a file descriptor into an array buffer,
+   * but here it simply logs the invocation for debugging purposes.
+   */
+  fd_read() {
+    debug('fd_read');
+  },
+  /**
+   * Handles the file descriptor write operation.
+   * This dummy implementation of the WASI 'fd_write' function is part of the wasi API:
+   * https://github.com/WebAssembly/WASI/blob/main/phases/snapshot/docs.md
+   * The environment expects this function, but it is not used in the current implementation.
+   * It is used to write data from WebAssembly memory to a file descriptor.
+   */
+  fd_write() {
+    debug('fd_write');
+  },
+  /**
+   * Perform a file seek operation on the given file descriptor to change its current position.
+   * The new position is calculated using the provided offset and whence values.
+   * Throws an error if the file descriptor is invalid or the operation cannot be performed.
+   *
+   * @param fd - The file descriptor of the file to perform the seek operation on.
+   * @param offset - The relative offset to apply, based on the whence value.
+   * @param whence - The reference point from which the offset should be calculated. One of SEEK_SET (start), SEEK_CUR (current), or SEEK_END (end).
+   * @returns The new position in the file after the seek operation has been performed.
+   */
+  fd_seek() {
+    debug('fd_seek');
+  },
+  /**
+   * This function is a dummy implementation of the 'fd_fdstat_get' function in the WebAssembly System Interface (WASI) API.
+   * Although not actually used in this context, it is present due to the environment's expectation of its existence.
+   * The 'fd_fdstat_get' function is typically responsible for obtaining file descriptor status information.
+   */
+  fd_fdstat_get() {
+    debug('fd_fdstat_get');
+  },
+  /**
+   * Sets the file descriptor flags for a given file descriptor.
+   * This function is a dummy implementation of the WASI API function 'fd_fdstat_set_flags'.
+   * It currently does not perform any operation but logs the function call with a debug instance.
+   * This is provided since the environment expects this function to be present.
+   */
+  fd_fdstat_set_flags() {
+    debug('fd_fdstat_set_flags');
+  },
+  /**
+   * Handles the `fd_prestat_get` function call for the dummy WebAssembly System Interface (WASI) implementation.
+   * This function is expected by the WASI environment, although it is not used in this implementation.
+   * The `fd_prestat_get` function retrieves pre-opened file descriptor properties.
+   *
+   * @returns A constant integer value indicating successful completion of the function call.
+   */
+  fd_prestat_get() {
+    debug('fd_prestat_get');
+    return 8;
+  },
+  /**
+   * Provides a dummy implementation for the `fd_prestat_dir_name` function, which is expected to be called by the WASI environment.
+   * This function is intended to retrieve the pre-opened directory's path associated with the given file descriptor. However, since it's a dummy implementation,
+   * it doesn't perform any actual operation and only logs the function call with the provided debug logger.
+   *
+   * @returns A constant number representing a dummy return value for the function call.
+   */
+  fd_prestat_dir_name() {
+    debug('fd_prestat_dir_name');
+    return 28;
+  },
+  /**
+   * Handles the opening of a file path within the WASI environment.
+   * This function is a dummy implementation required for compatibility with
+   * the WebAssembly System Interface (WASI) API, but it does not perform any
+   * actual file opening operation. It is mainly used for debugging purposes.
+   */
+  path_open() {
+    debug('path_open');
+  },
+  /**
+   * Retrieve file system information of the specified path.
+   * This function retrieves statistics, such as size and permissions, associated with the file or directory
+   * identified by the given path. In case of an error or non-existing path, appropriate debug logs will be generated.
+   *
+   * @returns An object containing file statistics like size, permissions, etc.
+   */
+  path_filestat_get() {
+    debug('path_filestat_get');
+  },
+  /**
+   * Terminate the process normally, performing the regular cleanup for terminating programs.
+   * The input 'status' represents the exit code and is used to indicate success or failure
+   * of the program execution. A zero value typically indicates successful execution,
+   * while non-zero values are treated as errors by the operating system.
+   *
+   * @param status - The exit code representing the success or failure of the program execution.
+   * @returns The exit status code.
+   */
+  proc_exit() {
+    debug('proc_exit');
+    return 52;
+  },
+  /**
+   * Generates a random number and returns it.
+   * This dummy implementation of 'random_get' method in the wasi API is expected by the environment.
+   * In this case, the function always returns 1 to maintain consistency with the environment's expectations.
+   *
+   * @returns A random number. In this implementation, always returns 1.
+   */
+  random_get() {
+    debug('random_get');
+    return 1;
+  },
+});

package/src/wasm/wasm/index.ts

@@ -0,0 +1 @@
+export { WasmModule, IWasmModule } from './wasm_module.js';

package/src/wasm/wasm/wasm_module.test.ts

@@ -0,0 +1,27 @@
+import { WasmModule } from './wasm_module.js';
+import { fileURLToPath } from 'url';
+import { readFile } from 'fs/promises';
+import { dirname } from 'path';
+
+/**
+ * Fetch a simple WASM.
+ */
+async function fetchCode() {
+  const __dirname = dirname(fileURLToPath(import.meta.url));
+  return await readFile(`${__dirname}/../fixtures/gcd.wasm`);
+}
+
+describe('simple wasm', () => {
+  let wasm!: WasmModule;
+
+  beforeAll(async () => {
+    wasm = new WasmModule(await fetchCode(), () => ({
+      /*no imports*/
+    }));
+    await wasm.init(1, 1, /* No init method */ null);
+  });
+
+  it('should call gcd with correct result', () => {
+    expect(wasm.call('gcd', 12312, 12123)).toBe(27);
+  });
+});

package/src/wasm/wasm/wasm_module.ts

@@ -0,0 +1,257 @@
+import { Buffer } from 'buffer';
+import { randomBytes } from 'crypto';
+import { createDebugLogger, DebugLogger } from '../../log/index.js';
+import { getEmptyWasiSdk } from './empty_wasi_sdk.js';
+import { MemoryFifo } from '../../fifo/index.js';
+
+/**
+ * The base shape of a WASM module providing low level memory and synchronous call access.
+ * Note that the Aztec codebase used to support asyncify but has fully moved away,
+ * using web workers if needed.
+ */
+export interface IWasmModule {
+  /**
+   * Low level call function, returns result as number.
+   * @param name - The function name.
+   * @param args - The function args.
+   * @returns The value as an integer (could be pointer).
+   */
+  call(name: string, ...args: any): number;
+
+  /**
+   * Get a slice of memory between two addresses.
+   * @param start - The start address.
+   * @param end - The end address.
+   * @returns A Uint8Array view of memory.
+   */
+  getMemorySlice(start: number, end: number): Uint8Array;
+
+  /**
+   * Write data into the heap.
+   * @param offset - The address to write data at.
+   * @param arr - The data to write.
+   */
+  writeMemory(offset: number, arr: Uint8Array): void;
+}
+
+/**
+ * WasmModule:
+ *  Helper over a webassembly module.
+ *  Assumes a few quirks.
+ *  1) the module expects wasi_snapshot_preview1 with the methods from getEmptyWasiSdk
+ *  2) of which the webassembly
+ *  we instantiate only uses random_get (update this if more WASI sdk methods are needed).
+ */
+export class WasmModule implements IWasmModule {
+  private memory!: WebAssembly.Memory;
+  private heap!: Uint8Array;
+  private instance?: WebAssembly.Instance;
+  private mutexQ = new MemoryFifo<boolean>();
+  private debug: DebugLogger;
+
+  /**
+   * Create a wasm module. Should be followed by await init();.
+   * @param module - The module as a WebAssembly.Module or a Buffer.
+   * @param importFn - Imports expected by the WASM.
+   * @param loggerName - Optional, for debug logging.
+   */
+  constructor(
+    private module: WebAssembly.Module | Buffer,
+    private importFn: (module: WasmModule) => any,
+    loggerName = 'wasm',
+  ) {
+    this.debug = createDebugLogger(loggerName);
+    this.mutexQ.put(true);
+  }
+
+  /**
+   * Return the wasm source.
+   * @returns The source.
+   */
+  public getModule(): WebAssembly.Module | Buffer {
+    return this.module;
+  }
+  /**
+   * Initialize this wasm module.
+   * @param wasmImportEnv - Linked to a module called "env". Functions implementations referenced from e.g. C++.
+   * @param initial - 30 pages by default. 30*2**16 \> 1mb stack size plus other overheads.
+   * @param initMethod - Defaults to calling '_initialize'.
+   * @param maximum - 8192 maximum by default. 512mb.
+   */
+  public async init(initial = 30, maximum = 8192, initMethod: string | null = '_initialize') {
+    this.debug(
+      `initial mem: ${initial} pages, ${(initial * 2 ** 16) / (1024 * 1024)}mb. max mem: ${maximum} pages, ${
+        (maximum * 2 ** 16) / (1024 * 1024)
+      }mb`,
+    );
+    this.memory = new WebAssembly.Memory({ initial, maximum });
+    // Create a view over the memory buffer.
+    // We do this once here, as webkit *seems* bugged out and actually shows this as new memory,
+    // thus displaying double. It's only worse if we create views on demand. I haven't established yet if
+    // the bug is also exasperating the termination on mobile due to "excessive memory usage". It could be
+    // that the OS is actually getting an incorrect reading in the same way the memory profiler does...
+    // The view will have to be recreated if the memory is grown. See getMemory().
+    this.heap = new Uint8Array(this.memory.buffer);
+
+    // We support the wasi 12 SDK, but only implement random_get
+    /* eslint-disable camelcase */
+    const importObj = {
+      wasi_snapshot_preview1: {
+        ...getEmptyWasiSdk(this.debug),
+        random_get: (arr: number, length: number) => {
+          arr = arr >>> 0;
+          const heap = this.getMemory();
+          const randomData = randomBytes(length);
+          for (let i = arr; i < arr + length; ++i) {
+            heap[i] = randomData[i - arr];
+          }
+        },
+      },
+      env: this.importFn(this),
+    };
+
+    if (this.module instanceof WebAssembly.Module) {
+      this.instance = await WebAssembly.instantiate(this.module, importObj);
+    } else {
+      const { instance } = await WebAssembly.instantiate(this.module, importObj);
+      this.instance = instance;
+    }
+
+    // Init all global/static data.
+    if (initMethod) {
+      this.call(initMethod);
+    }
+  }
+
+  /**
+   * The methods or objects exported by the WASM module.
+   * @returns An indexable object.
+   */
+  public exports(): any {
+    if (!this.instance) {
+      throw new Error('WasmModule: not initialized!');
+    }
+    return this.instance.exports;
+  }
+
+  /**
+   * Get the current logger.
+   * @returns Logging function.
+   */
+  public getLogger() {
+    return this.debug;
+  }
+
+  /**
+   * Add a logger.
+   * @param logger - Function to call when logging.
+   */
+  public addLogger(logger: DebugLogger) {
+    const oldDebug = this.debug;
+    this.debug = (...args: any[]) => {
+      logger(...args);
+      oldDebug(...args);
+    };
+  }
+
+  /**
+   * Calls into the WebAssembly.
+   * @param name - The method name.
+   * @param args - The arguments to the method.
+   * @returns The numeric method result.
+   */
+  public call(name: string, ...args: any): number {
+    if (!this.exports()[name]) {
+      throw new Error(`WASM function ${name} not found.`);
+    }
+    try {
+      // When returning values from the WASM, use >>> operator to convert
+      // signed representation to unsigned representation.
+      return this.exports()[name](...args) >>> 0;
+    } catch (err: any) {
+      const message = `WASM function ${name} aborted, error: ${err}\n${err.stack}`;
+      throw new Error(message);
+    }
+  }
+  /**
+   * Get the memory used by the WASM module.
+   * @returns A WebAssembly memory object.
+   */
+  public getRawMemory(): WebAssembly.Memory {
+    return this.memory;
+  }
+  /**
+   * Get the memory used by the WASM module, as a byte array.
+   * @returns A Uint8Array view of the WASM module memory.
+   */
+  public getMemory(): Uint8Array {
+    // If the memory is grown, our view over it will be lost. Recreate the view.
+    if (this.heap.length === 0) {
+      this.heap = new Uint8Array(this.memory.buffer);
+    }
+    return this.heap;
+  }
+
+  /**
+   * The memory size in bytes.
+   * @returns Number of bytes.
+   */
+  public memSize(): number {
+    return this.getMemory().length;
+  }
+
+  /**
+   * Get a slice of memory between two addresses.
+   * @param start - The start address.
+   * @param end - The end address.
+   * @returns A Uint8Array view of memory.
+   */
+  public getMemorySlice(start: number, end: number): Uint8Array {
+    return this.getMemory().slice(start, end);
+  }
+
+  /**
+   * Write data into the heap.
+   * @param offset - The address to write data at.
+   * @param arr - The data to write.
+   */
+  public writeMemory(offset: number, arr: Uint8Array) {
+    const mem = this.getMemory();
+    for (let i = 0; i < arr.length; i++) {
+      mem[i + offset] = arr[i];
+    }
+  }
+
+  /**
+   * Read WASM memory as a JS string.
+   * @param addr - The memory address.
+   * @returns A JS string.
+   */
+  public getMemoryAsString(addr: number) {
+    addr = addr >>> 0;
+    const m = this.getMemory();
+    let i = addr;
+    for (; m[i] !== 0; ++i);
+    return Buffer.from(m.slice(addr, i)).toString('ascii');
+  }
+
+  /**
+   * When calling the wasm, sometimes a caller will require exclusive access over a series of calls.
+   * E.g. When a result is written to address 0, one cannot have another caller writing to the same address via
+   * writeMemory before the result is read via sliceMemory.
+   * Acquire() gets a single token from a fifo. The caller must call release() to add the token back.
+   */
+  public async acquire() {
+    await this.mutexQ.get();
+  }
+
+  /**
+   * Release the mutex, letting another promise call acquire().
+   */
+  public release() {
+    if (this.mutexQ.length() !== 0) {
+      throw new Error('Release called but not acquired.');
+    }
+    this.mutexQ.put(true);
+  }
+}

package/src/wasm/worker/browser/index.ts

@@ -0,0 +1,2 @@
+export * from './web_worker.js';
+export * from './web_data_store.js';

package/src/wasm/worker/browser/start_web_module.ts

@@ -0,0 +1,23 @@
+import { DispatchMsg, TransportServer, WorkerListener } from '../../../transport/index.js';
+import { WasmModule } from '../../wasm/wasm_module.js';
+
+/**
+ * Start the transport server corresponding to this module.
+ * @param module - The WasmModule to host.
+ */
+export function startWebModule(module: WasmModule) {
+  const dispatch = async ({ fn, args }: DispatchMsg) => {
+    if (fn === '__destroyWorker__') {
+      transportServer.stop();
+      return;
+    }
+    if (!(module as any)[fn]) {
+      throw new Error(`dispatch error, function not found: ${fn}`);
+    }
+    return await (module as any)[fn](...args);
+  };
+  const transportListener = new WorkerListener(self);
+  const transportServer = new TransportServer<DispatchMsg>(transportListener, dispatch);
+  module.addLogger((...args: any[]) => transportServer.broadcast({ fn: 'emit', args: ['log', ...args] }));
+  transportServer.start();
+}

package/src/wasm/worker/browser/web_data_store.ts

@@ -0,0 +1,37 @@
+import { DataStore } from '../data_store.js';
+import levelup, { LevelUp } from 'levelup';
+import memdown from 'memdown';
+
+/**
+ * Cache for data used by wasm module.
+ * Stores in a LevelUp database.
+ */
+export class WebDataStore implements DataStore {
+  private db: LevelUp;
+
+  constructor() {
+    // TODO: The whole point of this is to reduce memory load in the browser.
+    // Replace with leveljs so the data is stored in indexeddb and not in memory.
+    // Hack: Cast as any to work around package "broken" with node16 resolution
+    // See https://github.com/microsoft/TypeScript/issues/49160
+    this.db = levelup((memdown as any)());
+  }
+
+  /**
+   * Lookup a key.
+   * @param key - Key to lookup.
+   * @returns The buffer.
+   */
+  async get(key: string): Promise<Buffer | undefined> {
+    return await this.db.get(key).catch(() => {});
+  }
+
+  /**
+   * Alter a key.
+   * @param key - Key to alter.
+   * @param value - Buffer to store.
+   */
+  async set(key: string, value: Buffer): Promise<void> {
+    await this.db.put(key, value);
+  }
+}

package/src/wasm/worker/browser/web_worker.ts

@@ -0,0 +1,24 @@
+import { DispatchMsg, TransportClient, WorkerConnector, createDispatchProxy } from '../../../transport/index.js';
+import { WasmModule } from '../../wasm/wasm_module.js';
+import { WasmWorker } from '../wasm_worker.js';
+
+/**
+ * Instantiate a web worker.
+ * @param url - The URL.
+ * @param initialMem - Initial memory pages.
+ * @param maxMem - Maximum memory pages.
+ * @returns The worker.
+ */
+export async function createWebWorker(url: string, initialMem?: number, maxMem?: number): Promise<WasmWorker> {
+  const worker = new Worker(url);
+  const transportConnect = new WorkerConnector(worker);
+  const transportClient = new TransportClient<DispatchMsg>(transportConnect);
+  await transportClient.open();
+  const remoteModule = createDispatchProxy(WasmModule, transportClient) as WasmWorker;
+  remoteModule.destroyWorker = async () => {
+    await transportClient.request({ fn: '__destroyWorker__', args: [] });
+    transportClient.close();
+  };
+  await remoteModule.init(initialMem, maxMem);
+  return remoteModule;
+}

package/src/wasm/worker/data_store.ts

@@ -0,0 +1,7 @@
+/**
+ * Simple read/write interface for wasm modules.
+ */
+export interface DataStore {
+  get(key: string): Promise<Buffer | undefined>;
+  set(key: string, value: Buffer): Promise<void>;
+}

package/src/wasm/worker/index.ts

@@ -0,0 +1,2 @@
+export { WorkerPool } from './worker_pool.js';
+export { WasmWorker } from './wasm_worker.js';

package/src/wasm/worker/node/index.ts

@@ -0,0 +1,2 @@
+export * from './node_worker.js';
+export * from './node_data_store.js';

package/src/wasm/worker/node/node_data_store.ts

@@ -0,0 +1,36 @@
+import { DataStore } from '../data_store.js';
+import levelup, { LevelUp } from 'levelup';
+import leveldown from 'leveldown';
+import memdown from 'memdown';
+
+/**
+ * Cache for data used by wasm module.
+ */
+export class NodeDataStore implements DataStore {
+  private db: LevelUp;
+
+  // eslint-disable-next-line
+  constructor(path?: string) {
+    // Hack: Cast as any to work around packages "broken" with node16 resolution
+    // See https://github.com/microsoft/TypeScript/issues/49160
+    this.db = levelup(path ? (leveldown as any)(path) : (memdown as any)());
+  }
+
+  /**
+   * Get a value from our DB.
+   * @param key - The key to look up.
+   * @returns The value.
+   */
+  async get(key: string): Promise<Buffer | undefined> {
+    return await this.db.get(key).catch(() => {});
+  }
+
+  /**
+   * Set a value in our DB.
+   * @param key - The key to update.
+   * @param value - The value to set.
+   */
+  async set(key: string, value: Buffer): Promise<void> {
+    await this.db.put(key, value);
+  }
+}

package/src/wasm/worker/node/node_worker.ts

@@ -0,0 +1,21 @@
+import { Worker } from 'worker_threads';
+import { WasmModule } from '../../wasm/wasm_module.js';
+import { WasmWorker } from '../wasm_worker.js';
+import { DispatchMsg, NodeConnector, TransportClient, createDispatchProxy } from '../../../transport/index.js';
+
+/**
+ * Creates a node worker.
+ */
+export async function createNodeWorker(filepath: string, initialMem?: number, maxMem?: number): Promise<WasmWorker> {
+  const worker = new Worker(filepath);
+  const transportConnect = new NodeConnector(worker);
+  const transportClient = new TransportClient<DispatchMsg>(transportConnect);
+  await transportClient.open();
+  const remoteModule = createDispatchProxy(WasmModule, transportClient) as WasmWorker;
+  remoteModule.destroyWorker = async () => {
+    await transportClient.request({ fn: '__destroyWorker__', args: [] });
+    transportClient.close();
+  };
+  await remoteModule.init(initialMem, maxMem);
+  return remoteModule;
+}

package/src/wasm/worker/node/start_node_module.ts

@@ -0,0 +1,28 @@
+import { parentPort } from 'worker_threads';
+import { WasmModule } from '../../wasm/wasm_module.js';
+import { DispatchMsg, NodeListener, TransportServer } from '../../../transport/index.js';
+
+if (!parentPort) {
+  throw new Error('InvalidWorker');
+}
+
+/**
+ * Start the transport server corresponding to this module.
+ * @param module - The WasmModule to host.
+ */
+export function startNodeModule(module: WasmModule) {
+  const dispatch = async ({ fn, args }: DispatchMsg) => {
+    if (fn === '__destroyWorker__') {
+      transportServer.stop();
+      return;
+    }
+    if (!(module as any)[fn]) {
+      throw new Error(`dispatch error, function not found: ${fn}`);
+    }
+    return await (module as any)[fn](...args);
+  };
+  const transportListener = new NodeListener();
+  const transportServer = new TransportServer<DispatchMsg>(transportListener, dispatch);
+  module.addLogger((...args: any[]) => transportServer.broadcast({ fn: 'emit', args: ['log', ...args] }));
+  transportServer.start();
+}

package/src/wasm/worker/wasm_worker.ts

@@ -0,0 +1,7 @@
+import { Proxify } from '../../transport/index.js';
+import { WasmModule } from '../wasm/wasm_module.js';
+
+/**
+ * Represents either a WASM web worker, or node.js worker.
+ */
+export type WasmWorker = Proxify<WasmModule> & { destroyWorker(): void };