@aztec/foundation 0.0.1-commit.b655e406 → 0.0.1-commit.b8a057fa

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

@@ -22,6 +22,7 @@ import { InterruptError } from '../error/index.js';
  */
 export class InterruptibleSleep {
   private interrupts: Array<(shouldThrow: boolean) => void> = [];
+  private timeoutIds: NodeJS.Timeout[] = [];
 
   /**
    * Sleep for a specified amount of time in milliseconds.
@@ -38,9 +39,15 @@ export class InterruptibleSleep {
       this.interrupts.push(resolve);
     });
 
-    const timeoutPromise = new Promise<boolean>(resolve => setTimeout(() => resolve(false), ms));
+    let timeoutId: NodeJS.Timeout;
+    const timeoutPromise = new Promise<boolean>(resolve => {
+      timeoutId = setTimeout(() => resolve(false), ms);
+      this.timeoutIds.push(timeoutId);
+    });
     const shouldThrow = await Promise.race([interruptPromise, timeoutPromise]);
 
+    clearTimeout(timeoutId!);
+    this.timeoutIds = this.timeoutIds.filter(id => id !== timeoutId);
     this.interrupts = this.interrupts.filter(res => res !== interruptResolve);
 
     if (shouldThrow) {
@@ -58,6 +65,8 @@ export class InterruptibleSleep {
   public interrupt(sleepShouldThrow = false): void {
     this.interrupts.forEach(resolve => resolve(sleepShouldThrow));
     this.interrupts = [];
+    this.timeoutIds.forEach(id => clearTimeout(id));
+    this.timeoutIds = [];
   }
 }
 
@@ -73,3 +82,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/string/index.ts

@@ -58,7 +58,7 @@ export function urlJoin(...args: string[]): string {
       end--;
     }
 
-    if (start < end) {
+    if (start <= end) {
       processed.push(arg.slice(start, end + 1));
     }
   }

package/src/testing/files/index.ts

@@ -52,8 +52,7 @@ export function updateInlineTestData(targetFileFromRepoRoot: string, itemName: s
 /**
  * Updates the sample Prover.toml files in noir-projects/noir-protocol-circuits/crates/.
  * @remarks Requires AZTEC_GENERATE_TEST_DATA=1 to be set
- * To re-gen, run 'AZTEC_GENERATE_TEST_DATA=1 FAKE_PROOFS=1 yarn test:e2e full.test'
- * To re-gen public base only, run 'AZTEC_GENERATE_TEST_DATA=1 yarn workspace @aztec/prover-client test orchestrator_public_functions'
+ * To re-gen, run 'AZTEC_GENERATE_TEST_DATA=1 FAKE_PROOFS=1 yarn test:e2e e2e_prover/full.test'
  */
 export function updateProtocolCircuitSampleInputs(circuitName: string, value: string) {
   const logger = createConsoleLogger('aztec:testing:test_data');
@@ -63,7 +62,7 @@ export function updateProtocolCircuitSampleInputs(circuitName: string, value: st
   logger(`Updated test data in ${targetFile} for ${circuitName}`);
 }
 
-function getPathToFile(targetFileFromRepoRoot: string) {
+export function getPathToFile(targetFileFromRepoRoot: string) {
   const repoRoot = resolve(dirname(fileURLToPath(import.meta.url)), '../../../../../');
   if (!existsSync(join(repoRoot, 'CODEOWNERS'))) {
     throw new Error(`Path to repo root is incorrect (got ${repoRoot})`);

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,52 @@ export class TestDateProvider extends DateProvider {
     this.offset = timeMs - Date.now();
     this.logger.warn(`Time set to ${new Date(timeMs).toISOString()}`, { offset: this.offset, timeMs });
   }
+
+  /** Resets the time back to real time (offset = 0). */
+  public reset() {
+    this.offset = 0;
+    this.logger.warn('Time reset to real time');
+  }
+
+  /** 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/timer/index.ts

@@ -1,4 +1,4 @@
 export * from './date.js';
 export { elapsed, elapsedSync } from './elapsed.js';
-export { TimeoutTask, executeTimeout, timeoutPromise } from './timeout.js';
+export { TimeoutTask, execWithSignal, executeTimeout, timeoutPromise } from './timeout.js';
 export { Timer } from './timer.js';

package/src/timer/timeout.ts

@@ -92,6 +92,37 @@ export async function executeTimeout<T>(
   return await task.exec();
 }
 
+function defaultSignalError(signal: AbortSignal) {
+  if (signal.reason instanceof Error) {
+    return signal.reason;
+  }
+
+  return new Error(signal.reason ? String(signal.reason) : 'Operation aborted');
+}
+
+/** Executes a function until it completes or the given signal aborts. */
+export async function execWithSignal<T>(
+  fn: (signal: AbortSignal) => Promise<T>,
+  signal: AbortSignal,
+  errorFn: (signal: AbortSignal) => Error = defaultSignalError,
+) {
+  if (signal.aborted) {
+    throw errorFn(signal);
+  }
+
+  const abortPromise = promiseWithResolvers<never>();
+  const abortListener = () => abortPromise.reject(errorFn(signal));
+  signal.addEventListener('abort', abortListener, { once: true });
+
+  try {
+    return await Promise.race<T>([fn(signal), abortPromise.promise]);
+  } finally {
+    if (abortListener) {
+      signal.removeEventListener('abort', abortListener);
+    }
+  }
+}
+
 /** Returns a promise that rejects after the given timeout */
 export function timeoutPromise(timeoutMs: number, errorMessage?: string) {
   const promise = promiseWithResolvers<never>();

package/src/transport/dispatch/create_dispatch_proxy.ts

@@ -2,7 +2,17 @@ import { EventEmitter } from 'events';
 
 import { type TransferDescriptor, isTransferDescriptor } from '../interface/transferable.js';
 import type { TransportClient } from '../transport_client.js';
-import type { DispatchMsg } from './create_dispatch_fn.js';
+
+/**
+ * Represents a message object for dispatching function calls.
+ * Contains the function name ('fn') and an array of arguments ('args') required to call the target method.
+ */
+export interface DispatchMsg {
+  /** Name of the target method to be called. */
+  fn: string;
+  /** An array of arguments to be passed to the target method. */
+  args: any[];
+}
 
 /**
  * FilterOutAttributes type filters out all non-method properties of an object, leaving only the attributes

package/src/transport/index.ts

@@ -1,4 +1,3 @@
-export * from './dispatch/create_dispatch_fn.js';
 export * from './dispatch/create_dispatch_proxy.js';
 export * from './dispatch/messages.js';
 export * from './interface/connector.js';

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_root.ts

@@ -0,0 +1,65 @@
+import { shaMerkleHash } from './hasher.js';
+
+export const computeBalancedShaRoot = (leaves: Buffer[]) => computeBalancedMerkleTreeRoot(leaves);
+
+/**
+ * Computes the Merkle root with the provided leaves **synchronously**.
+ * This method uses a synchronous hash function (defaults to `sha256Trunc`).
+ *
+ * @throws If the number of leaves is not a power of two.
+ */
+export function computeBalancedMerkleTreeRoot(leaves: Buffer[], hasher = shaMerkleHash): Buffer {
+  const height = getTreeHeight(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(l, r);
+    }
+    nodes = nodes.slice(0, j);
+  }
+
+  return nodes[0];
+}
+
+/**
+ * Computes the Merkle root with the provided leaves **asynchronously**.
+ * This method uses an asynchronous hash function (defaults to `poseidon2Hash`).
+ *
+ * @throws If the number of leaves is not a power of two.
+ */
+export async function computeBalancedMerkleTreeRootAsync(
+  leaves: Buffer[],
+  hasher: (left: Buffer, right: Buffer) => Promise<Buffer<ArrayBuffer>>,
+): Promise<Buffer> {
+  const height = getTreeHeight(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] = await hasher(l, r);
+    }
+    nodes = nodes.slice(0, j);
+  }
+
+  return nodes[0];
+}
+
+function getTreeHeight(leaves: 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');
+  }
+
+  return height;
+}

package/src/trees/hasher.ts

@@ -1,3 +1,6 @@
+import { poseidon2HashWithSeparator } from '../crypto/poseidon/index.js';
+import { sha256Trunc } from '../crypto/sha256/index.js';
+
 /**
  * Defines hasher interface used by Merkle trees.
  */
@@ -37,3 +40,12 @@ 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>;
+
+/** Creates a poseidon2 merkle hasher with the given domain separator. */
+export const makePoseidonMerkleHash =
+  (separator: number): AsyncHasher['hash'] =>
+  async (left: Buffer, right: Buffer) =>
+    (await poseidon2HashWithSeparator([left, right], separator)).toBuffer() as Buffer<ArrayBuffer>;

package/src/trees/index.ts

@@ -1,5 +1,4 @@
-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';
@@ -9,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/indexed_merkle_tree_calculator.ts

@@ -40,7 +40,11 @@ export class IndexedMerkleTreeCalculator<T extends IndexedTreeLeafPreimage, N ex
     }
     const sorted = values
       .map((v, i) => ({ value: v, index: i }))
-      .sort((a, b) => Number(toBigIntBE(b.value) - toBigIntBE(a.value)));
+      .sort((a, b): -1 | 0 | 1 => {
+        const aBigInt = toBigIntBE(a.value);
+        const bBigInt = toBigIntBE(b.value);
+        return aBigInt < bBigInt ? 1 : aBigInt > bBigInt ? -1 : 0;
+      });
     const indexedLeaves = sorted.map((item, i) => ({
       leaf: this.factory.fromBuffer(
         Buffer.concat([

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';

package/src/trees/merkle_tree_calculator.ts

@@ -1,8 +1,7 @@
-import { pedersenHash, sha256Trunc } from '@aztec/foundation/crypto';
-
-import type { AsyncHasher } from './hasher.js';
 import { MerkleTree } from './merkle_tree.js';
 
+type MerkleHashFn = (left: Buffer, right: Buffer) => Promise<Buffer<ArrayBuffer>>;
+
 /**
  * Merkle tree calculator.
  */
@@ -10,17 +9,12 @@ export class MerkleTreeCalculator {
   private constructor(
     private height: number,
     private zeroHashes: Buffer[],
-    private hasher: AsyncHasher['hash'],
+    private hasher: MerkleHashFn,
   ) {
     this.hasher = hasher;
   }
 
-  static async create(
-    height: number,
-    zeroLeaf: Buffer = Buffer.alloc(32),
-    hasher = async (left: Buffer, right: Buffer) =>
-      (await pedersenHash([left, right])).toBuffer() as Buffer<ArrayBuffer>,
-  ) {
+  static async create(height: number, zeroLeaf: Buffer = Buffer.alloc(32), hasher: MerkleHashFn) {
     const zeroHashes = [zeroLeaf];
     for (let i = 0; i < height; i++) {
       zeroHashes.push((await hasher(zeroHashes[i], zeroHashes[i])) as Buffer<ArrayBuffer>);
@@ -69,37 +63,4 @@ 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/sibling_path.ts

@@ -1,6 +1,7 @@
+import { z } from 'zod';
+
 import { makeTuple } from '../array/array.js';
-import { pedersenHash } from '../crypto/index.js';
-import { Fr } from '../fields/index.js';
+import { Fr } from '../curves/bn254/index.js';
 import { schemas } from '../schemas/index.js';
 import {
   type Tuple,
@@ -43,10 +44,14 @@ export class SiblingPath<N extends number> {
   }
 
   static schemaFor<N extends number>(size: N) {
-    return schemas.Buffer.transform(b => SiblingPath.fromBuffer(b) as SiblingPath<N>).refine(
-      path => path.pathSize === size,
-      path => ({ message: `Expected sibling path size ${size} but got ${path.pathSize}` }),
-    );
+    return schemas.Buffer.transform(b => SiblingPath.fromBuffer(b) as SiblingPath<N>).superRefine((path, ctx) => {
+      if (path.pathSize !== size) {
+        ctx.addIssue({
+          code: z.ZodIssueCode.custom,
+          message: `Expected sibling path size ${size} but got ${path.pathSize}`,
+        });
+      }
+    });
   }
 
   toJSON() {
@@ -172,7 +177,7 @@ export async function computeRootFromSiblingPath(
   leaf: Buffer,
   siblingPath: Buffer[],
   index: number,
-  hasher = async (left: Buffer, right: Buffer) => (await pedersenHash([left, right])).toBuffer(),
+  hasher: (left: Buffer, right: Buffer) => Promise<Buffer>,
 ) {
   let result = leaf;
   for (const sibling of siblingPath) {

package/src/trees/unbalanced_merkle_tree_calculator.ts

@@ -1,17 +1,7 @@
-import { sha256Trunc } from '../crypto/index.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),
-  hasher?: Hasher['hash'],
-): Buffer {
-  const calculator = UnbalancedMerkleTreeCalculator.create(leaves, valueToCompress, hasher);
-  return calculator.getRoot();
-}
-
 interface TreeNode {
   value: Buffer;
   leafIndex?: number;
@@ -32,12 +22,9 @@ export class UnbalancedMerkleTreeCalculator {
   public constructor(
     private readonly leaves: Buffer[],
     private readonly valueToCompress: Buffer,
+    private readonly emptyRoot: Buffer,
     private readonly hasher: Hasher['hash'],
   ) {
-    if (leaves.length === 0) {
-      throw Error('Cannot create a compressed unbalanced tree with 0 leaves.');
-    }
-
     this.store = new UnbalancedTreeStore(leaves.length);
     this.buildTree();
   }
@@ -45,9 +32,10 @@ export class UnbalancedMerkleTreeCalculator {
   static create(
     leaves: Buffer[],
     valueToCompress = Buffer.alloc(0),
-    hasher = (left: Buffer, right: Buffer) => sha256Trunc(Buffer.concat([left, right])) as Buffer<ArrayBuffer>,
+    emptyRoot = Buffer.alloc(32),
+    hasher = shaMerkleHash,
   ) {
-    return new UnbalancedMerkleTreeCalculator(leaves, valueToCompress, hasher);
+    return new UnbalancedMerkleTreeCalculator(leaves, valueToCompress, emptyRoot, hasher);
   }
 
   /**
@@ -113,8 +101,8 @@ export class UnbalancedMerkleTreeCalculator {
     // Start with the leaves that are not compressed.
     let toProcess = this.leafLocations.filter((_, i) => !this.leaves[i].equals(this.valueToCompress));
     if (!toProcess.length) {
-      // All leaves are compressed. Set 0 to the root.
-      this.store.setNode({ level: 0, index: 0 }, { value: Buffer.alloc(32) });
+      // All leaves are compressed. Set empty root to the root.
+      this.store.setNode({ level: 0, index: 0 }, { value: this.emptyRoot });
       return;
     }
 
@@ -138,10 +126,7 @@ export class UnbalancedMerkleTreeCalculator {
         const shouldShiftUp = !sibling || sibling.value.equals(this.valueToCompress);
         if (shouldShiftUp) {
           // The node becomes the parent if the sibling is a compressed leaf.
-          const isLeaf = this.shiftNodeUp(location, parentLocation);
-          if (!isLeaf) {
-            this.shiftChildrenUp(location);
-          }
+          this.shiftNodeUp(location, parentLocation);
         } else {
           // Hash the value with the (right) sibling and update the parent node.
           const node = this.store.getNode(location)!;
@@ -157,38 +142,35 @@ export class UnbalancedMerkleTreeCalculator {
     }
   }
 
-  private shiftNodeUp(fromLocation: TreeNodeLocation, toLocation: TreeNodeLocation): boolean {
-    const node = this.store.getNode(fromLocation)!;
-
-    this.store.setNode(toLocation, node);
-
-    const isLeaf = node.leafIndex !== undefined;
-    if (isLeaf) {
-      // Update the location if the node is a leaf.
-      this.leafLocations[node.leafIndex!] = toLocation;
+  private shiftNodeUp(from: TreeNodeLocation, to: TreeNodeLocation) {
+    // Collect all nodes that need to shift.
+    const nodesToShift = this.collectNodesToShift(from, to);
+    // Move all nodes to their new locations.
+    for (const { node, newLocation } of nodesToShift) {
+      this.store.setNode(newLocation, node);
+      if (node.leafIndex !== undefined) {
+        this.leafLocations[node.leafIndex] = newLocation;
+      }
     }
-
-    return isLeaf;
   }
 
-  private shiftChildrenUp(parent: TreeNodeLocation) {
-    const [left, right] = this.store.getChildLocations(parent);
-
-    const level = parent.level;
-    const groupSize = 2 ** level;
-    const computeNewLocation = (index: number) => ({
-      level,
-      index: Math.floor(index / (groupSize * 2)) * groupSize + (index % groupSize),
-    });
+  private collectNodesToShift(from: TreeNodeLocation, to: TreeNodeLocation) {
+    const node = this.store.getNode(from);
+    if (!node) {
+      return [];
+    }
 
-    const isLeftLeaf = this.shiftNodeUp(left, computeNewLocation(left.index));
-    const isRightLeaf = this.shiftNodeUp(right, computeNewLocation(right.index));
+    let result = [{ node, newLocation: to }];
 
-    if (!isLeftLeaf) {
-      this.shiftChildrenUp(left);
-    }
-    if (!isRightLeaf) {
-      this.shiftChildrenUp(right);
+    // If not a leaf, collect its children.
+    if (node.leafIndex === undefined) {
+      const [leftChild, rightChild] = this.store.getChildLocations(from);
+      const [leftChildNew, rightChildNew] = this.store.getChildLocations(to);
+      result = result
+        .concat(this.collectNodesToShift(leftChild, leftChildNew))
+        .concat(this.collectNodesToShift(rightChild, rightChildNew));
     }
+
+    return result;
   }
 }