@aztec/foundation 0.69.1 → 0.70.0

package/src/trees/unbalanced_merkle_tree.ts

@@ -0,0 +1,103 @@
+import { padArrayEnd } from '@aztec/foundation/collection';
+import { sha256Trunc } from '@aztec/foundation/crypto';
+
+/**
+ * Computes the merkle root for 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.
+ */
+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);
+    }
+  }
+
+  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);
+      } 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];
+      }
+    }
+    // reset the layers
+    thisLayer = nextLayer;
+    nextLayer = [];
+  }
+  // return the root
+  return thisLayer[0];
+}
+
+function getMaxBalancedTreeDepth(numLeaves: number) {
+  return Math.floor(Math.log2(numLeaves));
+}
+
+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 = getMaxBalancedTreeDepth(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);
+}

package/src/trees/unbalanced_tree_store.ts

@@ -0,0 +1,102 @@
+import { findLeafLevelAndIndex } from './unbalanced_merkle_tree.js';
+
+export interface TreeNodeLocation {
+  level: number;
+  index: number;
+}
+
+interface TreeNode<T> {
+  value: T;
+  location: TreeNodeLocation;
+}
+
+export class UnbalancedTreeStore<T> {
+  #nodeMapping: Map<string, TreeNode<T>> = new Map();
+  readonly #numLeaves: number;
+
+  constructor(numLeaves: number) {
+    this.#numLeaves = numLeaves;
+  }
+
+  setLeaf(leafIndex: number, value: T): TreeNodeLocation {
+    if (leafIndex >= this.#numLeaves) {
+      throw new Error(`Expected at most ${this.#numLeaves} leaves. Received a leaf at index ${leafIndex}.`);
+    }
+
+    const { level, indexAtLevel } = findLeafLevelAndIndex(this.#numLeaves, leafIndex);
+    const location = {
+      level,
+      index: indexAtLevel,
+    };
+    this.#nodeMapping.set(this.#getKey(location), {
+      location,
+      value,
+    });
+    return location;
+  }
+
+  setNode({ level, index }: TreeNodeLocation, value: T) {
+    const location = {
+      level,
+      index,
+    };
+    this.#nodeMapping.set(this.#getKey(location), {
+      location,
+      value,
+    });
+  }
+
+  getParentLocation({ level, index }: TreeNodeLocation): TreeNodeLocation {
+    if (level === 0) {
+      throw new Error('Tree root does not have a parent.');
+    }
+
+    return { level: level - 1, index: Math.floor(index / 2) };
+  }
+
+  getSiblingLocation({ level, index }: TreeNodeLocation): TreeNodeLocation {
+    if (level === 0) {
+      throw new Error('Tree root does not have a sibling.');
+    }
+
+    return { level, index: index % 2 ? index - 1 : index + 1 };
+  }
+
+  getChildLocations({ level, index }: TreeNodeLocation): [TreeNodeLocation, TreeNodeLocation] {
+    const left = { level: level + 1, index: index * 2 };
+    const right = { level: level + 1, index: index * 2 + 1 };
+    return [left, right];
+  }
+
+  getLeaf(leafIndex: number) {
+    const { level, indexAtLevel } = findLeafLevelAndIndex(this.#numLeaves, leafIndex);
+    const location = {
+      level,
+      index: indexAtLevel,
+    };
+    return this.getNode(location);
+  }
+
+  getNode(location: TreeNodeLocation): T | undefined {
+    return this.#nodeMapping.get(this.#getKey(location))?.value;
+  }
+
+  getParent(location: TreeNodeLocation): T | undefined {
+    const parentLocation = this.getParentLocation(location);
+    return this.getNode(parentLocation);
+  }
+
+  getSibling(location: TreeNodeLocation): T | undefined {
+    const siblingLocation = this.getSiblingLocation(location);
+    return this.getNode(siblingLocation);
+  }
+
+  getChildren(location: TreeNodeLocation): [T | undefined, T | undefined] {
+    const [left, right] = this.getChildLocations(location);
+    return [this.getNode(left), this.getNode(right)];
+  }
+
+  #getKey(location: TreeNodeLocation) {
+    return `${location.level}-${location.index}`;
+  }
+}

package/dest/trees/unbalanced_merkle_root.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"unbalanced_merkle_root.d.ts","sourceRoot":"","sources":["../../src/trees/unbalanced_merkle_root.ts"],"names":[],"mappings":";;AACA,OAAO,EAAE,WAAW,EAAE,MAAM,0BAA0B,CAAC;AAEvD;;;;;;;GAOG;AACH,wBAAgB,2BAA2B,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,SAAS,CAAC,EAAE,MAAM,EAAE,MAAM,qBAAc,GAAG,MAAM,CAwC9G"}

package/src/trees/unbalanced_merkle_root.ts

@@ -1,52 +0,0 @@
-import { padArrayEnd } from '@aztec/foundation/collection';
-import { sha256Trunc } from '@aztec/foundation/crypto';
-
-/**
- * Computes the merkle root for an unbalanced tree.
- *
- * @dev Adapted from proving-state.ts -> findMergeLevel and 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.
- */
-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);
-    }
-  }
-
-  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);
-      } 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];
-      }
-    }
-    // reset the layers
-    thisLayer = nextLayer;
-    nextLayer = [];
-  }
-  // return the root
-  return thisLayer[0];
-}