@aztec/foundation 1.2.0 → 2.0.0-nightly.20250813

package/src/timer/timeout.ts

@@ -8,23 +8,15 @@ import { TimeoutError } from '../error/index.js';
  * @typeparam T - The return type of the asynchronous function to be executed.
  */
 export class TimeoutTask<T> {
-  private interruptPromise!: Promise<any>;
-  private interrupt = () => {};
-  private interrupted = false;
   private totalTime = 0;
+  private timeoutPromise: Promise<never> | undefined;
 
   constructor(
     private fn: (signal: AbortSignal) => Promise<T>,
     private timeout: number,
-    errorFn: () => any,
-  ) {
-    this.interruptPromise = new Promise<T>((_, reject) => {
-      this.interrupt = () => {
-        this.interrupted = true;
-        reject(errorFn());
-      };
-    });
-  }
+    private errorFn: () => Error,
+    private onAbort?: () => void,
+  ) {}
 
   /**
    * Executes the given function with a specified timeout.
@@ -35,22 +27,22 @@ export class TimeoutTask<T> {
    * @throws An error with a message indicating the function was interrupted due to exceeding the specified timeout.
    */
   public async exec() {
-    const interruptTimeout = setTimeout(this.interrupt, this.timeout);
-    const controller = new AbortController();
-    try {
-      const start = Date.now();
-      const result = await Promise.race<T>([this.fn(controller.signal), this.interruptPromise]);
-      this.totalTime = Date.now() - start;
-      return result;
-    } catch (err) {
-      if (this.interrupted) {
-        controller.abort();
-      }
+    const signal = AbortSignal.timeout(this.timeout);
+    this.timeoutPromise = new Promise<never>((_, reject) => {
+      signal!.addEventListener(
+        'abort',
+        () => {
+          this.onAbort?.();
+          reject(this.errorFn());
+        },
+        { once: true },
+      );
+    });
 
-      throw err;
-    } finally {
-      clearTimeout(interruptTimeout);
-    }
+    const start = Date.now();
+    const result = await Promise.race<T>([this.fn(signal), this.timeoutPromise]);
+    this.totalTime = Date.now() - start;
+    return result;
   }
 
   /**
@@ -61,7 +53,7 @@ export class TimeoutTask<T> {
    * @returns The interrupt promise associated with the task.
    */
   public getInterruptPromise() {
-    return this.interruptPromise;
+    return this.timeoutPromise;
   }
 
   /**
@@ -75,15 +67,26 @@ export class TimeoutTask<T> {
   }
 }
 
+/**
+ * Executes a function with a timeout.
+ * @param fn - The function to execute, accepts AbortSignal and returns a Promise.
+ * @param timeout - The maximum time in milliseconds to allow the function to run.
+ * @param errorOrFnName - Optional function name or a function that returns an Error to throw if the timeout is reached.
+ * @param onAbort - Optional callback to execute when the task is aborted.
+ *
+ * @returns A Promise that resolves with the result of the function fn if it completes within the timeout
+ *
+ * */
 export async function executeTimeout<T>(
   fn: (signal: AbortSignal) => Promise<T>,
   timeout: number,
-  errorOrFnName?: string | (() => any),
+  errorOrFnName?: string | (() => Error),
+  onAbort?: () => void,
 ) {
   const errorFn =
     typeof errorOrFnName === 'function'
       ? errorOrFnName
       : () => new TimeoutError(`Timeout running ${errorOrFnName ?? 'function'} after ${timeout}ms.`);
-  const task = new TimeoutTask(fn, timeout, errorFn);
+  const task = new TimeoutTask(fn, timeout, errorFn, onAbort);
   return await task.exec();
 }

package/src/trees/merkle_tree_calculator.ts

@@ -1,4 +1,4 @@
-import { pedersenHash } from '@aztec/foundation/crypto';
+import { pedersenHash, sha256Trunc } from '@aztec/foundation/crypto';
 
 import type { AsyncHasher } from './hasher.js';
 import { MerkleTree } from './merkle_tree.js';
@@ -69,4 +69,37 @@ export class MerkleTreeCalculator {
 
     return leaves[0];
   }
+
+  /**
+   * Computes the Merkle root with the provided leaves **synchronously**.
+   *
+   * This method uses a synchronous hash function (defaults to `sha256Trunc`) and **does not** allow for padding.
+   * If the number of leaves is not a power of two, it throws an error.
+   * This contrasts with the above non-static async method `computeTreeRoot`, which can handle any number of leaves by
+   * padding with zero hashes.
+   */
+  static computeTreeRootSync(leaves: Buffer[], hasher = sha256Trunc): Buffer {
+    if (leaves.length === 0) {
+      throw new Error('Cannot compute a Merkle root with no leaves');
+    }
+
+    const height = Math.log2(leaves.length);
+    if (!Number.isInteger(height)) {
+      throw new Error('Cannot compute a Merkle root with a non-power-of-two number of leaves');
+    }
+
+    let nodes = leaves.slice();
+
+    for (let i = 0; i < height; ++i) {
+      let j = 0;
+      for (; j < nodes.length / 2; ++j) {
+        const l = nodes[j * 2];
+        const r = nodes[j * 2 + 1];
+        nodes[j] = hasher(Buffer.concat([l, r]));
+      }
+      nodes = nodes.slice(0, j);
+    }
+
+    return nodes[0];
+  }
 }

package/src/trees/unbalanced_merkle_tree.ts

@@ -1,61 +1,69 @@
-import { padArrayEnd } from '../collection/array.js';
 import { sha256Trunc } from '../crypto/sha256/index.js';
+import { MerkleTreeCalculator } from './merkle_tree_calculator.js';
 
 /**
- * Computes the merkle root for an unbalanced tree.
+ * Computes the Merkle root of an unbalanced tree.
  *
- * @dev Adapted from unbalanced_tree.ts.
- * Calculates the tree upwards layer by layer until we reach the root.
- * The L1 calculation instead computes the tree from right to left (slightly cheaper gas).
- * TODO: A more thorough investigation of which method is cheaper, then use that method everywhere.
+ * Unlike a balanced Merkle tree, which requires the number of leaves to be a power of two, an unbalanced tree can have
+ * any number of leaves.
+ *
+ * The tree is constructed by iteratively extracting the smallest power-of-two-sized subtrees from **right to left**.
+ * For each such subtree, it computes the subtree root and then combines all subtree roots (again from right to left)
+ * into a single root using the provided hash function.
+ *
+ * Note: We need the final tree to be as shallow as possible, to minimize the size of the sibling path required to prove
+ * membership of a leaf. Therefor, the computation proceeds from right to left - smaller subtrees must always be
+ * combined before being merged with a larger sibling on their left.
+ *
+ * For example, consider an unbalanced tree made of three subtrees of sizes 2, 4, and 8. If we combine the size-2 and
+ * size-4 subtrees first (producing a subtree of depth 3), and then merge it with the size-8 subtree (also depth 3), the
+ * resulting tree has a maximum depth of 4.
+ *
+ * But if we instead combine the size-4 and size-8 subtrees first (depth 4), and then merge with the size-2 subtree
+ * (depth 1), the final tree has a depth of 5.
  */
-export function computeUnbalancedMerkleRoot(leaves: Buffer[], emptyLeaf?: Buffer, hasher = sha256Trunc): Buffer {
-  // Pad leaves to 2
-  if (leaves.length < 2) {
-    if (emptyLeaf === undefined) {
-      throw new Error('Cannot compute a Merkle root with less than 2 leaves');
-    } else {
-      leaves = padArrayEnd(leaves, emptyLeaf, 2);
-    }
+export function computeUnbalancedMerkleTreeRoot(leaves: Buffer[], hasher = sha256Trunc): Buffer {
+  if (!leaves.length) {
+    throw new Error('Cannot compute a Merkle root with no leaves');
   }
 
-  const depth = Math.ceil(Math.log2(leaves.length));
-  let [layerWidth, nodeToShift] =
-    leaves.length & 1 ? [leaves.length - 1, leaves[leaves.length - 1]] : [leaves.length, Buffer.alloc(0)];
-  // Allocate this layer's leaves and init the next layer up
-  let thisLayer = leaves.slice(0, layerWidth);
-  let nextLayer = [];
-  for (let i = 0; i < depth; i++) {
-    for (let j = 0; j < layerWidth; j += 2) {
-      // Store the hash of each pair one layer up
-      nextLayer[j / 2] = hasher(Buffer.concat([thisLayer[j], thisLayer[j + 1]]));
-    }
-    layerWidth /= 2;
-    if (layerWidth & 1) {
-      if (nodeToShift.length) {
-        // If the next layer has odd length, and we have a node that needs to be shifted up, add it here
-        nextLayer.push(nodeToShift);
-        layerWidth += 1;
-        nodeToShift = Buffer.alloc(0);
+  if (leaves.length === 1) {
+    return leaves[0];
+  }
+
+  let numRemainingLeaves = leaves.length;
+  let subtreeSize = 1;
+  let root: Buffer | undefined;
+  while (numRemainingLeaves > 1) {
+    if ((numRemainingLeaves & subtreeSize) !== 0) {
+      const subtreeLeaves = leaves.slice(numRemainingLeaves - subtreeSize, numRemainingLeaves);
+      const subtreeRoot = MerkleTreeCalculator.computeTreeRootSync(subtreeLeaves, hasher);
+      if (!root) {
+        root = subtreeRoot;
       } else {
-        // If we don't have a node waiting to be shifted, store the next layer's final node to be shifted
-        layerWidth -= 1;
-        nodeToShift = nextLayer[layerWidth];
+        root = hasher(Buffer.concat([subtreeRoot, root]));
       }
+
+      numRemainingLeaves -= subtreeSize;
     }
-    // reset the layers
-    thisLayer = nextLayer;
-    nextLayer = [];
+
+    subtreeSize *= 2;
   }
-  // return the root
-  return thisLayer[0];
+
+  return root!;
 }
 
-function getMaxBalancedTreeDepth(numLeaves: number) {
+/// Get the depth of the maximum balanced tree that can be created with the given number of leaves. The subtree will be
+/// the left most subtree of the wonky tree with a total of `numLeaves` leaves.
+///
+/// Note: All the leaves may not be used to form the tree. For example, if there are 5 leaves, the maximum depth is 2,
+/// only 4 leaves are used to form a balanced tree.
+function getMaxBalancedSubtreeDepth(numLeaves: number) {
   return Math.floor(Math.log2(numLeaves));
 }
 
-function getMaxUnbalancedTreeDepth(numLeaves: number) {
+/// Get the maximum depth of an unbalanced tree that can be created with the given number of leaves.
+export function getMaxUnbalancedTreeDepth(numLeaves: number) {
   return Math.ceil(Math.log2(numLeaves));
 }
 
@@ -72,7 +80,7 @@ function findPosition(
   }
 
   // The largest balanced tree that can be created with the given number of leaves.
-  const maxBalancedTreeDepth = getMaxBalancedTreeDepth(numLeaves);
+  const maxBalancedTreeDepth = getMaxBalancedSubtreeDepth(numLeaves);
   const numBalancedLeaves = 2 ** maxBalancedTreeDepth;
   const numRemainingLeaves = numLeaves - numBalancedLeaves;
 

package/src/trees/unbalanced_merkle_tree_calculator.ts

@@ -1,8 +1,8 @@
-import { serializeToBuffer } from '@aztec/foundation/serialize';
+import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
 import type { AsyncHasher } from '@aztec/foundation/trees';
 import { SiblingPath } from '@aztec/foundation/trees';
 
-import { pedersenHash } from '../crypto/pedersen/index.js';
+import { sha256Trunc } from '../crypto/index.js';
 
 const indexToKeyHash = (level: number, index: bigint) => `${level}:${index}`;
 
@@ -10,6 +10,7 @@ const indexToKeyHash = (level: number, index: bigint) => `${level}:${index}`;
  * An ephemeral unbalanced Merkle tree implementation.
  * Follows the rollup implementation which greedily hashes pairs of nodes up the tree.
  * Remaining rightmost nodes are shifted up until they can be paired.
+ * If there is only one leaf, the root is the leaf.
  */
 export class UnbalancedMerkleTreeCalculator {
   // This map stores index and depth -> value
@@ -27,8 +28,8 @@ export class UnbalancedMerkleTreeCalculator {
 
   static create(
     height: number,
-    hasher = async (left: Buffer, right: Buffer) =>
-      (await pedersenHash([left, right])).toBuffer() as Buffer<ArrayBuffer>,
+    hasher = (left: Buffer, right: Buffer) =>
+      Promise.resolve(sha256Trunc(Buffer.concat([left, right])) as Buffer<ArrayBuffer>),
   ) {
     return new UnbalancedMerkleTreeCalculator(height, hasher);
   }
@@ -47,7 +48,11 @@ export class UnbalancedMerkleTreeCalculator {
    * @returns A sibling path for the element.
    * Note: The sibling path is an array of sibling hashes, with the lowest hash (leaf hash) first, and the highest hash last.
    */
-  public getSiblingPath<N extends number>(value: bigint): Promise<SiblingPath<N>> {
+  public getSiblingPath<N extends number>(value: Bufferable): Promise<SiblingPath<N>> {
+    if (this.size === 1n) {
+      return Promise.resolve(new SiblingPath<N>(0 as N, []));
+    }
+
     const path: Buffer[] = [];
     const [depth, _index] = this.valueCache[serializeToBuffer(value).toString('hex')].split(':');
     let level = parseInt(depth, 10);
@@ -72,8 +77,17 @@ export class UnbalancedMerkleTreeCalculator {
     if (this.size != BigInt(0)) {
       throw Error(`Can't re-append to an unbalanced tree. Current has ${this.size} leaves.`);
     }
-    const root = await this.batchInsert(leaves);
-    this.root = root;
+    if (leaves.length === 0) {
+      throw Error(`Can't append 0 leaves to an unbalanced tree.`);
+    }
+
+    if (leaves.length === 1) {
+      this.root = leaves[0];
+    } else {
+      this.root = await this.batchInsert(leaves);
+    }
+
+    this.size = BigInt(leaves.length);
 
     return Promise.resolve();
   }
@@ -119,7 +133,7 @@ export class UnbalancedMerkleTreeCalculator {
       thisLayer = nextLayer;
       nextLayer = [];
     }
-    this.size += BigInt(_leaves.length);
+
     // return the root
     return thisLayer[0];
   }

package/src/types/index.ts

@@ -31,9 +31,11 @@ export function unfreeze<T>(obj: T): Writeable<T> {
  * }
  */
 export interface TypedEventEmitter<TEventMap extends { [key in keyof TEventMap]: (...args: any[]) => void }> {
+  once<K extends keyof TEventMap>(event: K, listener: TEventMap[K]): this;
   on<K extends keyof TEventMap>(event: K, listener: TEventMap[K]): this;
   off<K extends keyof TEventMap>(event: K, listener: TEventMap[K]): this;
   emit<K extends keyof TEventMap>(event: K, ...args: Parameters<TEventMap[K]>): boolean;
   removeListener<K extends keyof TEventMap>(event: K, listener: TEventMap[K]): this;
+  removeAllListeners<K extends keyof TEventMap>(event: K): this;
   // Can add other EventEmitter methods if needed, like once(), listenerCount(), etc.
 }