@aztec/foundation 0.0.1-commit.5476d83 → 0.0.1-commit.5914bae

package/src/serialize/buffer_reader.ts

@@ -1,3 +1,4 @@
+import { toBigIntBE } from '../bigint-buffer/index.js';
 import type { Tuple } from './types.js';
 
 /**
@@ -130,6 +131,20 @@ export class BufferReader {
     return result;
   }
 
+  /**
+   * Reads a 256-bit signed integer (two's complement) from the buffer at the current index position.
+   * Updates the index position by 32 bytes after reading the number.
+   *
+   * @returns The read 256 bit signed value as a bigint.
+   */
+  public readInt256(): bigint {
+    this.#rangeCheck(32);
+    const unsigned = toBigIntBE(this.buffer.subarray(this.index, this.index + 32));
+    this.index += 32;
+    const signBit = 1n << 255n;
+    return unsigned >= signBit ? unsigned - (1n << 256n) : unsigned;
+  }
+
   /** Alias for readUInt256 */
   public readBigInt(): bigint {
     return this.readUInt256();
@@ -224,15 +239,22 @@ export class BufferReader {
    * deserializing each one using the 'fromBuffer' method of 'itemDeserializer'.
    *
    * @param itemDeserializer - Object with 'fromBuffer' method to deserialize vector elements.
+   * @param maxSize - Optional maximum allowed size for the vector. If the size exceeds this, an error is thrown.
    * @returns An array of deserialized elements of type T.
    */
-  public readVector<T>(itemDeserializer: {
-    /**
-     * A method to deserialize data from a buffer.
-     */
-    fromBuffer: (reader: BufferReader) => T;
-  }): T[] {
+  public readVector<T>(
+    itemDeserializer: {
+      /**
+       * A method to deserialize data from a buffer.
+       */
+      fromBuffer: (reader: BufferReader) => T;
+    },
+    maxSize?: number,
+  ): T[] {
     const size = this.readNumber();
+    if (maxSize !== undefined && size > maxSize) {
+      throw new Error(`Vector size ${size} exceeds maximum allowed ${maxSize}`);
+    }
     const result = new Array<T>(size);
     for (let i = 0; i < size; i++) {
       result[i] = itemDeserializer.fromBuffer(this);
@@ -264,16 +286,39 @@ export class BufferReader {
   }
 
   /**
-   * Read an array of a fixed size with elements of type T from the buffer.
-   * The 'itemDeserializer' object should have a 'fromBuffer' method that takes a BufferReader instance as input,
-   * and returns an instance of the desired deserialized data type T.
-   * This method will call the 'fromBuffer' method for each element in the array and return the resulting array.
+   * Read an array from the buffer using lazy allocation (new Array + loop).
+   * Safe for use with untrusted sizes — does not pre-allocate memory proportional to size.
    *
-   * @param size - The fixed number of elements in the array.
+   * @param size - The number of elements to read.
    * @param itemDeserializer - An object with a 'fromBuffer' method to deserialize individual elements of type T.
    * @returns An array of instances of type T.
    */
-  public readArray<T, N extends number>(
+  public readArray<T>(
+    size: number,
+    itemDeserializer: {
+      /**
+       * A function for deserializing data from a BufferReader instance.
+       */
+      fromBuffer: (reader: BufferReader) => T;
+    },
+  ): T[] {
+    const result = new Array<T>(size);
+    for (let i = 0; i < size; i++) {
+      result[i] = itemDeserializer.fromBuffer(this);
+    }
+    return result;
+  }
+
+  /**
+   * Read a fixed-size tuple from the buffer using dense allocation (Array.from).
+   * Only use with compile-time constant sizes — the size parameter MUST NOT come from untrusted input
+   * as Array.from pre-allocates memory proportional to size.
+   *
+   * @param size - The fixed number of elements (must be a compile-time constant).
+   * @param itemDeserializer - An object with a 'fromBuffer' method to deserialize individual elements of type T.
+   * @returns A densely-allocated tuple of instances of type T.
+   */
+  public readTuple<T, N extends number>(
     size: N,
     itemDeserializer: {
       /**
@@ -344,10 +389,11 @@ export class BufferReader {
    * The method first reads the size of the string, then reads the corresponding
    * number of bytes from the buffer and converts them to a string.
    *
+   * @param maxSize - Optional maximum allowed size for the string buffer. If the size exceeds this, an error is thrown.
    * @returns The read string from the buffer.
    */
-  public readString(): string {
-    return this.readBuffer().toString();
+  public readString(maxSize?: number): string {
+    return this.readBuffer(maxSize).toString();
   }
 
   /**
@@ -356,10 +402,14 @@ export class BufferReader {
    * a Buffer with that size containing the bytes. Useful for reading variable-length
    * binary data encoded as (size, data) format.
    *
+   * @param maxSize - Optional maximum allowed size for the buffer. If the size exceeds this, an error is thrown.
    * @returns A Buffer containing the read bytes.
    */
-  public readBuffer(): Buffer {
+  public readBuffer(maxSize?: number): Buffer {
     const size = this.readNumber();
+    if (maxSize !== undefined && size > maxSize) {
+      throw new Error(`Buffer size ${size} exceeds maximum allowed ${maxSize}`);
+    }
     this.#rangeCheck(size);
     return this.readBytes(size);
   }

package/src/serialize/field_reader.ts

@@ -1,4 +1,4 @@
-import { Fq, type Fr } from '../fields/fields.js';
+import { Fq, type Fr } from '../curves/bn254/field.js';
 import type { Tuple } from './types.js';
 
 /**
@@ -169,12 +169,30 @@ export class FieldReader {
    * @param itemDeserializer - An object with a 'fromFields' method to deserialize individual elements of type T.
    * @returns An array of instances of type T.
    */
-  public readArray<T, N extends number>(
+  /**
+   * Read an array from the field array using lazy allocation (new Array + loop).
+   * Safe for use with untrusted sizes.
+   */
+  public readArray<T>(
+    size: number,
+    itemDeserializer: {
+      fromFields: (reader: FieldReader) => T;
+    },
+  ): T[] {
+    const result = new Array<T>(size);
+    for (let i = 0; i < size; i++) {
+      result[i] = itemDeserializer.fromFields(this);
+    }
+    return result;
+  }
+
+  /**
+   * Read a fixed-size tuple from the field array using dense allocation (Array.from).
+   * Only use with compile-time constant sizes — the size parameter MUST NOT come from untrusted input.
+   */
+  public readTuple<T, N extends number>(
     size: N,
     itemDeserializer: {
-      /**
-       * A function for deserializing data from a FieldReader instance.
-       */
       fromFields: (reader: FieldReader) => T;
     },
   ): Tuple<T, N> {

package/src/serialize/free_funcs.ts

@@ -1,5 +1,5 @@
 import { toBufferBE } from '../bigint-buffer/index.js';
-import { Fr } from '../fields/fields.js';
+import { Fr } from '../curves/bn254/field.js';
 import type { Tuple } from './types.js';
 
 /**

package/src/serialize/serialize.ts

@@ -1,5 +1,5 @@
 import { toBigIntBE, toBufferBE } from '../bigint-buffer/index.js';
-import { Fr } from '../fields/fields.js';
+import { Fr } from '../curves/bn254/field.js';
 import { numToUInt32BE } from './free_funcs.js';
 
 /**
@@ -269,6 +269,22 @@ export function serializeBigInt(n: bigint, width = 32) {
   return toBufferBE(n, width);
 }
 
+/**
+ * Serialize a signed BigInt value into a Buffer of specified width using two's complement.
+ * @param n - The signed BigInt value to be serialized.
+ * @param width - The width (in bytes) of the output Buffer, optional with default value 32.
+ * @returns A Buffer containing the serialized signed BigInt value in big-endian format.
+ */
+export function serializeSignedBigInt(n: bigint, width = 32) {
+  const widthBits = BigInt(width * 8);
+  const max = 1n << (widthBits - 1n);
+  if (n < -max || n >= max) {
+    throw new Error(`Signed BigInt ${n.toString()} does not fit into ${width} bytes`);
+  }
+  const unsigned = n < 0n ? (1n << widthBits) + n : n;
+  return toBufferBE(unsigned, width);
+}
+
 /**
  * Deserialize a big integer from a buffer, given an offset and width.
  * Reads the specified number of bytes from the buffer starting at the offset, converts it to a big integer, and returns the deserialized result along with the number of bytes read (advanced).
@@ -282,6 +298,22 @@ export function deserializeBigInt(buf: Buffer, offset = 0, width = 32) {
   return { elem: toBigIntBE(buf.subarray(offset, offset + width)), adv: width };
 }
 
+/**
+ * Deserialize a signed BigInt from a buffer (two's complement).
+ * @param buf - The buffer containing the signed big integer to be deserialized.
+ * @param offset - The position in the buffer where the integer starts. Defaults to 0.
+ * @param width - The number of bytes to read from the buffer for the integer. Defaults to 32.
+ * @returns An object containing the deserialized signed bigint value ('elem') and bytes advanced ('adv').
+ */
+export function deserializeSignedBigInt(buf: Buffer, offset = 0, width = 32) {
+  const { elem, adv } = deserializeBigInt(buf, offset, width);
+  const widthBits = BigInt(width * 8);
+  const signBit = 1n << (widthBits - 1n);
+  const fullRange = 1n << widthBits;
+  const signed = elem >= signBit ? elem - fullRange : elem;
+  return { elem: signed, adv };
+}
+
 /**
  * Serializes a Date object into a Buffer containing its timestamp as a big integer value.
  * The resulting Buffer has a fixed width of 8 bytes, representing a 64-bit big-endian integer.

package/src/sleep/index.ts

@@ -73,3 +73,9 @@ export class InterruptibleSleep {
 export function sleep<T>(ms: number, returnValue?: T): Promise<T> {
   return new Promise(resolve => setTimeout(() => resolve(returnValue as T), ms));
 }
+
+/** Sleeps until the target date */
+export function sleepUntil<T>(target: Date, now: Date, returnValue?: T): Promise<T> {
+  const ms = target.getTime() - now.getTime();
+  return sleep(ms, returnValue);
+}

package/src/timer/date.ts

@@ -9,6 +9,10 @@ export class DateProvider {
   public nowInSeconds(): number {
     return Math.floor(this.now() / 1000);
   }
+
+  public nowAsDate(): Date {
+    return new Date(this.now());
+  }
 }
 
 /** Returns current datetime and allows to override it. */
@@ -27,4 +31,46 @@ export class TestDateProvider extends DateProvider {
     this.offset = timeMs - Date.now();
     this.logger.warn(`Time set to ${new Date(timeMs).toISOString()}`, { offset: this.offset, timeMs });
   }
+
+  /** Advances the time by the given number of seconds. */
+  public advanceTime(seconds: number) {
+    this.offset += seconds * 1000;
+  }
+}
+
+/**
+ * A date provider for tests that only advances time via explicit advanceTime() calls.
+ * Unlike TestDateProvider, this does NOT track real time progression - time is completely
+ * frozen until explicitly advanced. This eliminates flakiness from tests taking
+ * varying amounts of real time to execute.
+ */
+export class ManualDateProvider extends DateProvider {
+  private currentTimeMs: number;
+
+  /**
+   * @param initialTimeMs - Initial time in milliseconds. Defaults to a round timestamp for easy visualization.
+   */
+  constructor(initialTimeMs: number = Date.UTC(2025, 0, 1, 0, 0, 0)) {
+    super();
+    this.currentTimeMs = initialTimeMs;
+  }
+
+  public override now(): number {
+    return this.currentTimeMs;
+  }
+
+  /** Sets the current time to the given timestamp in milliseconds. */
+  public setTime(timeMs: number) {
+    this.currentTimeMs = timeMs;
+  }
+
+  /** Advances the time by the given number of seconds. */
+  public advanceTime(seconds: number) {
+    this.currentTimeMs += seconds * 1000;
+  }
+
+  /** Advances the time by the given number of milliseconds. */
+  public advanceTimeMs(ms: number) {
+    this.currentTimeMs += ms;
+  }
 }

package/src/transport/transport_client.ts

@@ -91,7 +91,7 @@ export class TransportClient<Payload> extends EventEmitter {
     }
     const msgId = this.msgId++;
     const msg = { msgId, payload };
-    log.debug(format(`->`, msg));
+    log.trace(format(`->`, msg));
     return new Promise<any>((resolve, reject) => {
       this.pendingRequests.push({ resolve, reject, msgId });
       this.socket!.send(msg, transfer).catch(reject);
@@ -111,7 +111,7 @@ export class TransportClient<Payload> extends EventEmitter {
       this.close();
       return;
     }
-    log.debug(format(`<-`, msg));
+    log.trace(format(`<-`, msg));
     if (isEventMessage(msg)) {
       this.emit('event_msg', msg.payload);
       return;

package/src/trees/{balanced_merkle_tree.ts → balanced_merkle_tree_root.ts}

@@ -1,21 +1,7 @@
-import { pedersenHash as pedersenHashArray, poseidon2Hash, sha256Trunc } from '@aztec/foundation/crypto';
-
-import type { AsyncHasher, Hasher } from './hasher.js';
-
-export const shaMerkleHash: Hasher['hash'] = (left: Buffer, right: Buffer) =>
-  sha256Trunc(Buffer.concat([left, right])) as Buffer<ArrayBuffer>;
-
-export const pedersenMerkleHash: AsyncHasher['hash'] = async (left: Buffer, right: Buffer) =>
-  (await pedersenHashArray([left, right])).toBuffer() as Buffer<ArrayBuffer>;
-
-export const poseidonMerkleHash: AsyncHasher['hash'] = async (left: Buffer, right: Buffer) =>
-  (await poseidon2Hash([left, right])).toBuffer() as Buffer<ArrayBuffer>;
+import { poseidonMerkleHash, shaMerkleHash } from './hasher.js';
 
 export const computeBalancedShaRoot = (leaves: Buffer[]) => computeBalancedMerkleTreeRoot(leaves);
 
-export const computeBalancedPedersenRoot = async (leaves: Buffer[]) =>
-  await computeBalancedMerkleTreeRootAsync(leaves, pedersenMerkleHash);
-
 export const computeBalancedPoseidonRoot = async (leaves: Buffer[]) =>
   await computeBalancedMerkleTreeRootAsync(leaves, poseidonMerkleHash);
 
@@ -44,7 +30,7 @@ export function computeBalancedMerkleTreeRoot(leaves: Buffer[], hasher = shaMerk
 
 /**
  * Computes the Merkle root with the provided leaves **asynchronously**.
- * This method uses an asynchronous hash function (defaults to `pedersenHash`).
+ * This method uses an asynchronous hash function (defaults to `poseidon2Hash`).
  *
  * @throws If the number of leaves is not a power of two.
  */

package/src/trees/hasher.ts

@@ -1,3 +1,6 @@
+import { poseidon2Hash } from '../crypto/poseidon/index.js';
+import { sha256Trunc } from '../crypto/sha256/index.js';
+
 /**
  * Defines hasher interface used by Merkle trees.
  */
@@ -37,3 +40,9 @@ export interface AsyncHasher {
    */
   hashInputs(inputs: Buffer[]): Promise<Buffer<ArrayBuffer>>;
 }
+
+export const shaMerkleHash: Hasher['hash'] = (left: Buffer, right: Buffer) =>
+  sha256Trunc(Buffer.concat([left, right])) as Buffer<ArrayBuffer>;
+
+export const poseidonMerkleHash: AsyncHasher['hash'] = async (left: Buffer, right: Buffer) =>
+  (await poseidon2Hash([left, right])).toBuffer() as Buffer<ArrayBuffer>;

package/src/trees/index.ts

@@ -1,6 +1,4 @@
-export * from './balanced_merkle_tree.js';
-export * from './unbalanced_merkle_tree.js';
-export * from './unbalanced_tree_store.js';
+export * from './balanced_merkle_tree_root.js';
 export * from './merkle_tree_calculator.js';
 export * from './merkle_tree.js';
 export * from './indexed_merkle_tree_calculator.js';
@@ -10,3 +8,5 @@ export * from './membership_witness.js';
 export * from './hasher.js';
 export * from './indexed_tree_leaf.js';
 export * from './unbalanced_merkle_tree_calculator.js';
+export * from './unbalanced_merkle_tree_root.js';
+export * from './unbalanced_tree_store.js';

package/src/trees/indexed_merkle_tree.ts

@@ -1,4 +1,4 @@
-import { Fr } from '../fields/index.js';
+import { Fr } from '../curves/bn254/index.js';
 import { assertLength } from '../serialize/index.js';
 import type { IndexedTreeLeafPreimage } from './indexed_tree_leaf.js';
 import { MembershipWitness } from './membership_witness.js';

package/src/trees/membership_witness.ts

@@ -1,6 +1,6 @@
 import { assertMemberLength } from '../array/array.js';
 import { toBigIntBE, toBufferBE } from '../bigint-buffer/index.js';
-import { Fr } from '../fields/fields.js';
+import { Fr } from '../curves/bn254/field.js';
 import { schemas } from '../schemas/schemas.js';
 import { BufferReader, type Tuple, serializeToBuffer } from '../serialize/index.js';
 import type { SiblingPath } from './sibling_path.js';
@@ -38,6 +38,14 @@ export class MembershipWitness<N extends number> {
     return [new Fr(this.leafIndex), ...this.siblingPath];
   }
 
+  /**
+   * Returns a representation of the membership witness as expected by intrinsic Noir deserialization.
+   */
+  public toNoirRepresentation(): (string | string[])[] {
+    // TODO(#12874): remove the stupid as string conversion by modifying ForeignCallOutput type in acvm.js
+    return [new Fr(this.leafIndex).toString() as string, this.siblingPath.map(fr => fr.toString()) as string[]];
+  }
+
   static schemaFor<N extends number>(size: N) {
     return schemas.Buffer.transform(b => MembershipWitness.fromBuffer(b, size));
   }
@@ -86,7 +94,7 @@ export class MembershipWitness<N extends number> {
   static fromBuffer<N extends number>(buffer: Buffer | BufferReader, size: N): MembershipWitness<N> {
     const reader = BufferReader.asReader(buffer);
     const leafIndex = toBigIntBE(reader.readBytes(32));
-    const siblingPath = reader.readArray(size, Fr);
+    const siblingPath = reader.readArray(size, Fr) as Tuple<Fr, N>;
     return new MembershipWitness(size, leafIndex, siblingPath);
   }
 
@@ -100,7 +108,7 @@ export class MembershipWitness<N extends number> {
       fromBuffer: (buffer: Buffer | BufferReader) => {
         const reader = BufferReader.asReader(buffer);
         const leafIndex = toBigIntBE(reader.readBytes(32));
-        const siblingPath = reader.readArray(size, Fr);
+        const siblingPath = reader.readArray(size, Fr) as Tuple<Fr, N>;
         return new MembershipWitness(size, leafIndex, siblingPath);
       },
     };

package/src/trees/merkle_tree_calculator.ts

@@ -1,4 +1,4 @@
-import { pedersenHash } from '@aztec/foundation/crypto';
+import { poseidon2Hash } from '@aztec/foundation/crypto/poseidon';
 
 import type { AsyncHasher } from './hasher.js';
 import { MerkleTree } from './merkle_tree.js';
@@ -19,7 +19,7 @@ export class MerkleTreeCalculator {
     height: number,
     zeroLeaf: Buffer = Buffer.alloc(32),
     hasher = async (left: Buffer, right: Buffer) =>
-      (await pedersenHash([left, right])).toBuffer() as Buffer<ArrayBuffer>,
+      (await poseidon2Hash([left, right])).toBuffer() as Buffer<ArrayBuffer>,
   ) {
     const zeroHashes = [zeroLeaf];
     for (let i = 0; i < height; i++) {

package/src/trees/sibling_path.ts

@@ -1,6 +1,6 @@
 import { makeTuple } from '../array/array.js';
-import { pedersenHash } from '../crypto/index.js';
-import { Fr } from '../fields/index.js';
+import { poseidon2Hash } from '../crypto/poseidon/index.js';
+import { Fr } from '../curves/bn254/index.js';
 import { schemas } from '../schemas/index.js';
 import {
   type Tuple,
@@ -172,7 +172,7 @@ export async function computeRootFromSiblingPath(
   leaf: Buffer,
   siblingPath: Buffer[],
   index: number,
-  hasher = async (left: Buffer, right: Buffer) => (await pedersenHash([left, right])).toBuffer(),
+  hasher = async (left: Buffer, right: Buffer) => (await poseidon2Hash([left, right])).toBuffer(),
 ) {
   let result = leaf;
   for (const sibling of siblingPath) {

package/src/trees/unbalanced_merkle_tree_calculator.ts

@@ -1,18 +1,7 @@
-import { shaMerkleHash } from './balanced_merkle_tree.js';
-import type { Hasher } from './hasher.js';
+import { type Hasher, shaMerkleHash } from './hasher.js';
 import { SiblingPath } from './sibling_path.js';
 import { type TreeNodeLocation, UnbalancedTreeStore } from './unbalanced_tree_store.js';
 
-export function computeCompressedUnbalancedMerkleTreeRoot(
-  leaves: Buffer[],
-  valueToCompress = Buffer.alloc(32),
-  emptyRoot = Buffer.alloc(32),
-  hasher?: Hasher['hash'],
-): Buffer {
-  const calculator = UnbalancedMerkleTreeCalculator.create(leaves, valueToCompress, emptyRoot, hasher);
-  return calculator.getRoot();
-}
-
 interface TreeNode {
   value: Buffer;
   leafIndex?: number;

package/src/trees/{unbalanced_merkle_tree.ts → unbalanced_merkle_tree_root.ts}

@@ -1,15 +1,14 @@
-import {
-  computeBalancedMerkleTreeRoot,
-  computeBalancedMerkleTreeRootAsync,
-  poseidonMerkleHash,
-  shaMerkleHash,
-} from './balanced_merkle_tree.js';
+import { computeBalancedMerkleTreeRoot, computeBalancedMerkleTreeRootAsync } from './balanced_merkle_tree_root.js';
+import { poseidonMerkleHash, shaMerkleHash } from './hasher.js';
+import { UnbalancedMerkleTreeCalculator } from './unbalanced_merkle_tree_calculator.js';
 
 export const computeUnbalancedShaRoot = (leaves: Buffer[]) => computeUnbalancedMerkleTreeRoot(leaves, shaMerkleHash);
 
 export const computeUnbalancedPoseidonRoot = async (leaves: Buffer[]) =>
   await computeUnbalancedMerkleTreeRootAsync(leaves, poseidonMerkleHash);
 
+export const computeWonkyShaRoot = (leaves: Buffer[]) => computeWonkyMerkleTreeRoot(leaves);
+
 /**
  * Computes the Merkle root of an unbalanced tree.
  *
@@ -101,59 +100,16 @@ export async function computeUnbalancedMerkleTreeRootAsync(
   return root!;
 }
 
-/// 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));
-}
-
-/// Get the maximum depth of an unbalanced tree that can be created with the given number of leaves.
-function getMaxUnbalancedTreeDepth(numLeaves: number) {
-  return Math.ceil(Math.log2(numLeaves));
-}
-
-function findPosition(
-  rootLevel: number,
-  leafLevel: number,
-  numLeaves: number,
-  indexOffset: number,
-  targetIndex: number,
-): { level: number; indexAtLevel: number } {
-  if (numLeaves <= 1) {
-    // Single leaf.
-    return { level: rootLevel, indexAtLevel: indexOffset };
-  }
-
-  // The largest balanced tree that can be created with the given number of leaves.
-  const maxBalancedTreeDepth = getMaxBalancedSubtreeDepth(numLeaves);
-  const numBalancedLeaves = 2 ** maxBalancedTreeDepth;
-  const numRemainingLeaves = numLeaves - numBalancedLeaves;
-
-  if (targetIndex < numBalancedLeaves) {
-    // Target is in the balanced tree.
-
-    // - If numRemainingLeaves is 0: this balanced tree is grown from the current root.
-    // - If numRemainingLeaves is not 0: the remaining leaves will form another tree, which will become the right child of the root.
-    //   And the balanced tree will be the left child of the root.
-    //   There will be an extra level between the root of the balanced tree and the current root.
-    const extraLevel = numRemainingLeaves ? 1 : 0;
-
-    return { level: rootLevel + maxBalancedTreeDepth + extraLevel, indexAtLevel: indexOffset + targetIndex };
-  } else {
-    // Target is in the right branch.
-    const rightBranchMaxLevel = getMaxUnbalancedTreeDepth(numRemainingLeaves);
-    const shiftedUp = leafLevel - rootLevel - rightBranchMaxLevel - 1;
-    const nextLeafLevel = leafLevel - shiftedUp;
-    const newIndexOffset = (indexOffset + numBalancedLeaves) >> shiftedUp;
-    const shiftedTargetIndex = targetIndex - numBalancedLeaves;
-    return findPosition(rootLevel + 1, nextLeafLevel, numRemainingLeaves, newIndexOffset, shiftedTargetIndex);
-  }
-}
-
-export function findLeafLevelAndIndex(numLeaves: number, leafIndex: number) {
-  const maxLevel = getMaxUnbalancedTreeDepth(numLeaves);
-  return findPosition(0, maxLevel, numLeaves, 0, leafIndex);
+// A **wonky** tree is a "compressed" unbalanced Merkle tree.
+// It is constructed in the same way as an unbalanced tree: by first creating the largest possible left subtree, with
+// the remaining leaves forming a right subtree that follows the same process recursively.
+// During construction, leaves equal to `valueToCompress` are skipped (compressed) and do not contribute to the tree.
+export function computeWonkyMerkleTreeRoot(
+  leaves: Buffer[],
+  valueToCompress = Buffer.alloc(32),
+  emptyRoot = Buffer.alloc(32),
+  hasher = shaMerkleHash,
+): Buffer {
+  const calculator = UnbalancedMerkleTreeCalculator.create(leaves, valueToCompress, emptyRoot, hasher);
+  return calculator.getRoot();
 }

package/src/trees/unbalanced_tree_store.ts

@@ -1,5 +1,3 @@
-import { findLeafLevelAndIndex } from './unbalanced_merkle_tree.js';
-
 export interface TreeNodeLocation {
   level: number;
   index: number;
@@ -104,3 +102,60 @@ export class UnbalancedTreeStore<T> {
     return `${location.level}-${location.index}`;
   }
 }
+
+/// 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 unbalanced 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));
+}
+
+/// Get the maximum depth of an unbalanced tree that can be created with the given number of leaves.
+function getMaxUnbalancedTreeDepth(numLeaves: number) {
+  return Math.ceil(Math.log2(numLeaves));
+}
+
+function findPosition(
+  rootLevel: number,
+  leafLevel: number,
+  numLeaves: number,
+  indexOffset: number,
+  targetIndex: number,
+): { level: number; indexAtLevel: number } {
+  if (numLeaves <= 1) {
+    // Single leaf.
+    return { level: rootLevel, indexAtLevel: indexOffset };
+  }
+
+  // The largest balanced tree that can be created with the given number of leaves.
+  const maxBalancedTreeDepth = getMaxBalancedSubtreeDepth(numLeaves);
+  const numBalancedLeaves = 2 ** maxBalancedTreeDepth;
+  const numRemainingLeaves = numLeaves - numBalancedLeaves;
+
+  if (targetIndex < numBalancedLeaves) {
+    // Target is in the balanced tree.
+
+    // - If numRemainingLeaves is 0: this balanced tree is grown from the current root.
+    // - If numRemainingLeaves is not 0: the remaining leaves will form another tree, which will become the right child of the root.
+    //   And the balanced tree will be the left child of the root.
+    //   There will be an extra level between the root of the balanced tree and the current root.
+    const extraLevel = numRemainingLeaves ? 1 : 0;
+
+    return { level: rootLevel + maxBalancedTreeDepth + extraLevel, indexAtLevel: indexOffset + targetIndex };
+  } else {
+    // Target is in the right branch.
+    const rightBranchMaxLevel = getMaxUnbalancedTreeDepth(numRemainingLeaves);
+    const shiftedUp = leafLevel - rootLevel - rightBranchMaxLevel - 1;
+    const nextLeafLevel = leafLevel - shiftedUp;
+    const newIndexOffset = (indexOffset + numBalancedLeaves) >> shiftedUp;
+    const shiftedTargetIndex = targetIndex - numBalancedLeaves;
+    return findPosition(rootLevel + 1, nextLeafLevel, numRemainingLeaves, newIndexOffset, shiftedTargetIndex);
+  }
+}
+
+export function findLeafLevelAndIndex(numLeaves: number, leafIndex: number) {
+  const maxLevel = getMaxUnbalancedTreeDepth(numLeaves);
+  return findPosition(0, maxLevel, numLeaves, 0, leafIndex);
+}