@aztec/foundation 0.0.1-fake-ceab37513c → 0.0.2-commit.217f559981

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/string/index.ts

@@ -1,3 +1,6 @@
+/** Parameterized hex string type for specific byte lengths */
+export type Hex<TByteLength extends number> = `0x${string}` & { readonly _length: TByteLength };
+
 export function hasHexPrefix(str: string): str is `0x${string}` {
   return str.startsWith('0x');
 }

package/src/testing/files/index.ts

@@ -63,7 +63,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/testing/formatting.ts

@@ -0,0 +1,3 @@
+export function toInlineStrArray(arr: { toString: () => string }[]): string {
+  return `[${arr.map(f => f.toString()).join(',')}]`;
+}

package/src/testing/index.ts

@@ -1,3 +1,4 @@
+export * from './formatting.js';
 export * from './snapshot_serializer.js';
 export * from './port_allocator.js';
 export * from './test_data.js';

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/trees/balanced_merkle_tree_root.ts

@@ -0,0 +1,68 @@
+import { poseidonMerkleHash, shaMerkleHash } from './hasher.js';
+
+export const computeBalancedShaRoot = (leaves: Buffer[]) => computeBalancedMerkleTreeRoot(leaves);
+
+export const computeBalancedPoseidonRoot = async (leaves: Buffer[]) =>
+  await computeBalancedMerkleTreeRootAsync(leaves, poseidonMerkleHash);
+
+/**
+ * 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 = poseidonMerkleHash,
+): 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 { 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,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/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));
   }

package/src/trees/merkle_tree_calculator.ts

@@ -1,4 +1,4 @@
-import { pedersenHash, sha256Trunc } 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++) {
@@ -69,37 +69,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,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,37 +1,41 @@
-import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize';
-import type { AsyncHasher } from '@aztec/foundation/trees';
-import { SiblingPath } from '@aztec/foundation/trees';
+import { type Hasher, shaMerkleHash } from './hasher.js';
+import { SiblingPath } from './sibling_path.js';
+import { type TreeNodeLocation, UnbalancedTreeStore } from './unbalanced_tree_store.js';
 
-import { sha256Trunc } from '../crypto/index.js';
-
-const indexToKeyHash = (level: number, index: bigint) => `${level}:${index}`;
+interface TreeNode {
+  value: Buffer;
+  leafIndex?: number;
+}
 
 /**
  * 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.
+ * The values that match the `valueToCompress` are skipped and the sibling of the compressed leaf 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
-  private cache: { [key: string]: Buffer } = {};
-  // This map stores value -> index and depth, since we have variable depth
-  private valueCache: { [key: string]: string } = {};
-  protected size: bigint = 0n;
-
-  root: Buffer = Buffer.alloc(32);
+  private store: UnbalancedTreeStore<TreeNode>;
+  private leafLocations: TreeNodeLocation[] = [];
 
   public constructor(
-    private maxDepth: number,
-    private hasher: AsyncHasher['hash'],
-  ) {}
+    private readonly leaves: Buffer[],
+    private readonly valueToCompress: Buffer,
+    private readonly emptyRoot: Buffer,
+    private readonly hasher: Hasher['hash'],
+  ) {
+    this.store = new UnbalancedTreeStore(leaves.length);
+    this.buildTree();
+  }
 
   static create(
-    height: number,
-    hasher = (left: Buffer, right: Buffer) =>
-      Promise.resolve(sha256Trunc(Buffer.concat([left, right])) as Buffer<ArrayBuffer>),
+    leaves: Buffer[],
+    valueToCompress = Buffer.alloc(0),
+    emptyRoot = Buffer.alloc(32),
+    hasher = shaMerkleHash,
   ) {
-    return new UnbalancedMerkleTreeCalculator(height, hasher);
+    return new UnbalancedMerkleTreeCalculator(leaves, valueToCompress, emptyRoot, hasher);
   }
 
   /**
@@ -39,108 +43,134 @@ export class UnbalancedMerkleTreeCalculator {
    * @returns The root of the tree.
    */
   public getRoot(): Buffer {
-    return this.root;
+    return this.store.getRoot()!.value;
   }
 
   /**
-   * Returns a sibling path for the element at the given index.
+   * Returns a sibling path for the element.
    * @param value - The value of the element.
    * @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: Bufferable): Promise<SiblingPath<N>> {
-    if (this.size === 1n) {
-      return Promise.resolve(new SiblingPath<N>(0 as N, []));
+  public getSiblingPath<N extends number>(value: Buffer): SiblingPath<N> {
+    const leafIndex = this.leaves.findIndex(leaf => leaf.equals(value));
+    if (leafIndex === -1) {
+      throw Error(`Leaf value ${value.toString('hex')} not found in tree.`);
     }
 
-    const path: Buffer[] = [];
-    const [depth, _index] = this.valueCache[serializeToBuffer(value).toString('hex')].split(':');
-    let level = parseInt(depth, 10);
-    let index = BigInt(_index);
-    while (level > 0) {
-      const isRight = index & 0x01n;
-      const key = indexToKeyHash(level, isRight ? index - 1n : index + 1n);
-      const sibling = this.cache[key];
-      path.push(sibling);
-      level -= 1;
-      index >>= 1n;
-    }
-    return Promise.resolve(new SiblingPath<N>(parseInt(depth, 10) as N, path));
+    return this.getSiblingPathByLeafIndex(leafIndex);
   }
 
   /**
-   * Appends the given leaves to the tree.
-   * @param leaves - The leaves to append.
-   * @returns Empty promise.
+   * Returns a sibling path for the leaf at the given index.
+   * @param leafIndex - The index of the leaf.
+   * @returns A sibling path for the leaf.
    */
-  public async appendLeaves(leaves: Buffer[]): Promise<void> {
-    if (this.size != BigInt(0)) {
-      throw Error(`Can't re-append to an unbalanced tree. Current has ${this.size} leaves.`);
+  public getSiblingPathByLeafIndex<N extends number>(leafIndex: number): SiblingPath<N> {
+    if (leafIndex >= this.leaves.length) {
+      throw Error(`Leaf index ${leafIndex} out of bounds. Tree has ${this.leaves.length} leaves.`);
     }
-    if (leaves.length === 0) {
-      throw Error(`Can't append 0 leaves to an unbalanced tree.`);
+
+    const leaf = this.leaves[leafIndex];
+    if (leaf.equals(this.valueToCompress)) {
+      throw Error(`Leaf at index ${leafIndex} has been compressed.`);
     }
 
-    if (leaves.length === 1) {
-      this.root = leaves[0];
-    } else {
-      this.root = await this.batchInsert(leaves);
+    const path: Buffer[] = [];
+    let location = this.leafLocations[leafIndex];
+    while (location.level > 0) {
+      const sibling = this.store.getSibling(location)!;
+      path.push(sibling.value);
+      location = this.store.getParentLocation(location);
     }
 
-    this.size = BigInt(leaves.length);
+    return new SiblingPath<N>(path.length as N, path);
+  }
 
-    return Promise.resolve();
+  public getLeafLocation(leafIndex: number) {
+    return this.leafLocations[leafIndex];
   }
 
   /**
-   * Calculates root while adding leaves and nodes to the cache.
-   * @param leaves - The leaves to append.
-   * @returns Resulting root of the tree.
+   * Adds leaves and nodes to the store. Updates the leafLocations.
+   * @param leaves - The leaves of the tree.
    */
-  private async batchInsert(_leaves: Buffer[]): Promise<Buffer> {
-    // If we have an even number of leaves, hash them all in pairs
-    // Otherwise, store the final leaf to be shifted up to the next odd sized level
-    let [layerWidth, nodeToShift] =
-      _leaves.length & 1
-        ? [_leaves.length - 1, serializeToBuffer(_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).map(l => serializeToBuffer(l));
-    let nextLayer = [];
-    // Store the bottom level leaves
-    thisLayer.forEach((leaf, i) => this.storeNode(leaf, this.maxDepth, BigInt(i)));
-    for (let i = 0; i < this.maxDepth; i++) {
-      for (let j = 0; j < layerWidth; j += 2) {
-        // Store the hash of each pair one layer up
-        nextLayer[j / 2] = await this.hasher(serializeToBuffer(thisLayer[j]), serializeToBuffer(thisLayer[j + 1]));
-        this.storeNode(nextLayer[j / 2], this.maxDepth - i - 1, BigInt(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(serializeToBuffer(nodeToShift));
-          this.storeNode(nodeToShift, this.maxDepth - i - 1, BigInt((layerWidth * 2) >> 1));
-          layerWidth += 1;
-          nodeToShift = Buffer.alloc(0);
+  private buildTree() {
+    this.leafLocations = this.leaves.map((value, i) => this.store.setLeaf(i, { value, leafIndex: i }));
+
+    // 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 empty root to the root.
+      this.store.setNode({ level: 0, index: 0 }, { value: this.emptyRoot });
+      return;
+    }
+
+    const level = toProcess[0].level;
+    for (let i = level; i > 0; i--) {
+      const toProcessNext = [];
+      for (const location of toProcess) {
+        if (location.level !== i) {
+          toProcessNext.push(location);
+          continue;
+        }
+
+        const parentLocation = this.store.getParentLocation(location);
+        if (this.store.getNode(parentLocation)) {
+          // Parent has been updated by its (left) sibling.
+          continue;
+        }
+
+        const sibling = this.store.getSibling(location);
+        // If sibling is undefined, all its children are compressed.
+        const shouldShiftUp = !sibling || sibling.value.equals(this.valueToCompress);
+        if (shouldShiftUp) {
+          // The node becomes the parent if the sibling is a compressed leaf.
+          this.shiftNodeUp(location, parentLocation);
         } 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];
+          // Hash the value with the (right) sibling and update the parent node.
+          const node = this.store.getNode(location)!;
+          const parentValue = this.hasher(node.value, sibling.value);
+          this.store.setNode(parentLocation, { value: parentValue });
         }
+
+        // Add the parent location to be processed next.
+        toProcessNext.push(parentLocation);
       }
-      // reset the layers
-      thisLayer = nextLayer;
-      nextLayer = [];
+
+      toProcess = toProcessNext;
     }
+  }
 
-    // return the root
-    return thisLayer[0];
+  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;
+      }
+    }
   }
 
-  private storeNode(value: Buffer, depth: number, index: bigint) {
-    const key = indexToKeyHash(depth, index);
-    this.cache[key] = value;
-    this.valueCache[value.toString('hex')] = key;
+  private collectNodesToShift(from: TreeNodeLocation, to: TreeNodeLocation) {
+    const node = this.store.getNode(from);
+    if (!node) {
+      return [];
+    }
+
+    let result = [{ node, newLocation: to }];
+
+    // 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;
   }
 }