fix(merkle-tree): make sibling path generic over its length (#647)

* feat: make sibling paths generic over their size

cleanup

fix: missed occurences

* fix: add merkle-tree to circuits

* fix: comments

* fix: yeeteth le todo

build_manifest.json

@@ -174,6 +174,7 @@
     "dependencies": [
       "barretenberg.js",
       "foundation",
+      "merkle-tree",
       "yarn-project-base"
     ]
   },

yarn-project/acir-simulator/src/client/private_execution.test.ts

@@ -186,7 +186,7 @@ describe('Private Execution test suite', () => {
           notes: await Promise.all(
             preimages.map(async (preimage, index) => ({
               preimage,
-              siblingPath: (await tree.getSiblingPath(BigInt(index), false)).data.map(buf => Fr.fromBuffer(buf)),
+              siblingPath: (await tree.getSiblingPath(BigInt(index), false)).toFieldArray(),
               index: BigInt(index),
             })),
           ),
@@ -261,7 +261,7 @@ describe('Private Execution test suite', () => {
           notes: await Promise.all(
             preimages.map(async (preimage, index) => ({
               preimage,
-              siblingPath: (await tree.getSiblingPath(BigInt(index), false)).data.map(buf => Fr.fromBuffer(buf)),
+              siblingPath: (await tree.getSiblingPath(BigInt(index), false)).toFieldArray(),
               index: BigInt(index),
             })),
           ),
@@ -419,7 +419,7 @@ describe('Private Execution test suite', () => {
         return Promise.resolve({
           message: preimage.toFieldArray(),
           index: 0n,
-          siblingPath: (await tree.getSiblingPath(0n, false)).data.map(buf => Fr.fromBuffer(buf)),
+          siblingPath: (await tree.getSiblingPath(0n, false)).toFieldArray(),
         });
       });
 

yarn-project/aztec-node/src/aztec-node/aztec-node.ts

@@ -1,6 +1,11 @@
 import { Archiver } from '@aztec/archiver';
 import { PrimitivesWasm } from '@aztec/barretenberg.js/wasm';
-import { CircuitsWasm } from '@aztec/circuits.js';
+import {
+  CONTRACT_TREE_HEIGHT,
+  CircuitsWasm,
+  L1_TO_L2_MESSAGES_TREE_HEIGHT,
+  PRIVATE_DATA_TREE_HEIGHT,
+} from '@aztec/circuits.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr } from '@aztec/foundation/fields';
 import { SiblingPath } from '@aztec/merkle-tree';
@@ -175,7 +180,7 @@ export class AztecNode {
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public getContractPath(leafIndex: bigint): Promise<SiblingPath> {
+  public getContractPath(leafIndex: bigint): Promise<SiblingPath<typeof CONTRACT_TREE_HEIGHT>> {
     return this.merkleTreeDB.getSiblingPath(MerkleTreeId.CONTRACT_TREE, leafIndex, false);
   }
 
@@ -184,7 +189,7 @@ export class AztecNode {
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public getDataTreePath(leafIndex: bigint): Promise<SiblingPath> {
+  public getDataTreePath(leafIndex: bigint): Promise<SiblingPath<typeof PRIVATE_DATA_TREE_HEIGHT>> {
     return this.merkleTreeDB.getSiblingPath(MerkleTreeId.PRIVATE_DATA_TREE, leafIndex, false);
   }
 
@@ -193,7 +198,7 @@ export class AztecNode {
    * @param leafIndex - Index of the leaf in the tree.
    * @returns The sibling path.
    */
-  public getL1ToL2MessagesTreePath(leafIndex: bigint): Promise<SiblingPath> {
+  public getL1ToL2MessagesTreePath(leafIndex: bigint): Promise<SiblingPath<typeof L1_TO_L2_MESSAGES_TREE_HEIGHT>> {
     return this.merkleTreeDB.getSiblingPath(MerkleTreeId.L1_TO_L2_MESSAGES_TREE, leafIndex, false);
   }
 

yarn-project/aztec-rpc/src/contract_tree/contract_tree.ts

@@ -263,10 +263,7 @@ export class ContractTree {
       this.contractMembershipWitness = new MembershipWitness<typeof CONTRACT_TREE_HEIGHT>(
         CONTRACT_TREE_HEIGHT,
         index,
-        assertLength(
-          siblingPath.data.map(x => Fr.fromBuffer(x)),
-          CONTRACT_TREE_HEIGHT,
-        ),
+        assertLength(siblingPath.toFieldArray(), CONTRACT_TREE_HEIGHT),
       );
     }
     return this.contractMembershipWitness;

yarn-project/aztec-rpc/src/simulator_oracle/index.ts

@@ -56,7 +56,7 @@ export class SimulatorOracle implements DBOracle {
           const path = await this.node.getDataTreePath(noteDao.index);
           return {
             preimage: noteDao.notePreimage.items,
-            siblingPath: path.data.map(buf => Fr.fromBuffer(buf)),
+            siblingPath: path.toFieldArray(),
             index: noteDao.index,
           };
         }),
@@ -103,7 +103,7 @@ export class SimulatorOracle implements DBOracle {
     const siblingPath = await this.node.getL1ToL2MessagesTreePath(index);
     return {
       message,
-      siblingPath: siblingPath.data.map(node => Fr.fromBuffer(node)),
+      siblingPath: siblingPath.toFieldArray(),
       index,
     };
   }

yarn-project/circuits.js/package.json

@@ -39,6 +39,7 @@
   "dependencies": {
     "@aztec/barretenberg.js": "workspace:^",
     "@aztec/foundation": "workspace:^",
+    "@aztec/merkle-tree": "workspace:^",
     "@msgpack/msgpack": "^3.0.0-beta2",
     "@types/lodash.camelcase": "^4.3.7",
     "@types/lodash.times": "^4.3.7",

yarn-project/circuits.js/src/structs/membership_witness.ts

@@ -1,4 +1,5 @@
 import { Fr } from '@aztec/foundation/fields';
+import { SiblingPath } from '@aztec/merkle-tree';
 import { assertMemberLength, range } from '../utils/jsUtils.js';
 import { serializeToBuffer } from '../utils/serialize.js';
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
@@ -72,4 +73,8 @@ export class MembershipWitness<N extends number> {
       siblingPath.map(x => Fr.fromBuffer(x)) as Tuple<Fr, N>,
     );
   }
+
+  static fromSiblingPath<N extends number>(leafIndex: bigint, siblingPath: SiblingPath<N>): MembershipWitness<N> {
+    return new MembershipWitness<N>(siblingPath.pathSize, leafIndex, siblingPath.toFieldArray() as Tuple<Fr, N>);
+  }
 }

yarn-project/circuits.js/tsconfig.json

@@ -11,6 +11,9 @@
     },
     {
       "path": "../foundation"
+    },
+    {
+      "path": "../merkle-tree"
     }
   ],
   "include": ["src"]

yarn-project/merkle-tree/src/interfaces/indexed_tree.ts

@@ -1,4 +1,4 @@
-import { LowLeafWitnessData } from '../index.js';
+import { LowLeafWitnessData, SiblingPath } from '../index.js';
 import { AppendOnlyTree } from './append_only_tree.js';
 
 /**
@@ -66,10 +66,13 @@ export interface IndexedTree extends AppendOnlyTree {
    * @param subtreeHeight - Height of the subtree.
    * @param includeUncommitted - If true, the uncommitted changes are included in the search.
    */
-  batchInsert(
+  batchInsert<TreeHeight extends number, SubtreeHeight extends number, SubtreeSiblingPathHeight extends number>(
     leaves: Buffer[],
-    treeHeight: number,
-    subtreeHeight: number,
+    treeHeight: TreeHeight,
+    subtreeHeight: SubtreeHeight,
     includeUncommitted: boolean,
-  ): Promise<[LowLeafWitnessData[], Buffer[]] | [undefined, Buffer[]]>;
+  ): Promise<
+    | [LowLeafWitnessData<TreeHeight>[], SiblingPath<SubtreeSiblingPathHeight>]
+    | [undefined, SiblingPath<SubtreeSiblingPathHeight>]
+  >;
 }

yarn-project/merkle-tree/src/interfaces/merkle_tree.ts

@@ -9,7 +9,7 @@ export interface SiblingPathSource {
    * @param index - The index of the leaf for which a sibling path is required.
    * @param includeUncommitted - Set to true to include uncommitted updates in the sibling path.
    */
-  getSiblingPath(index: bigint, includeUncommitted: boolean): Promise<SiblingPath>;
+  getSiblingPath<N extends number>(index: bigint, includeUncommitted: boolean): Promise<SiblingPath<N>>;
 }
 
 /**

yarn-project/merkle-tree/src/sibling_path/sibling_path.ts

@@ -1,5 +1,11 @@
-import { deserializeArrayFromVector, serializeBufferArrayToVector } from '@aztec/foundation/serialize';
+import {
+  Tuple,
+  assertLength,
+  deserializeArrayFromVector,
+  serializeBufferArrayToVector,
+} from '@aztec/foundation/serialize';
 import { Pedersen } from '../pedersen.js';
+import { Fr } from '@aztec/foundation/fields';
 
 /**
  * Contains functionality to compute and serialize/deserialize a sibling path.
@@ -11,34 +17,43 @@ import { Pedersen } from '../pedersen.js';
  *
  *      And the elements would be ordered as: [ leaf_at_index_2, node_at_level_2_index_0, node_at_level_1_index_1 ].
  */
-export class SiblingPath {
+export class SiblingPath<N extends number> {
+  private data: Tuple<Buffer, N>;
+
   /**
    * Returns sibling path hashed up from the a element.
    * @param size - The number of elements in a given path.
    * @param zeroElement - Value of the zero element.
    * @param pedersen - Implementation of a hasher interface using the Pedersen hash.
    * @returns A sibling path hashed up from a zero element.
    */
-  public static ZERO(size: number, zeroElement: Buffer, pedersen: Pedersen): SiblingPath {
+  public static ZERO<N extends number>(size: N, zeroElement: Buffer, pedersen: Pedersen): SiblingPath<N> {
     const bufs: Buffer[] = [];
     let current = zeroElement;
     for (let i = 0; i < size; ++i) {
       bufs.push(current);
       current = pedersen.compress(current, current);
     }
-    return new SiblingPath(bufs);
+    return new SiblingPath(size, bufs);
   }
 
   /**
    * Constructor.
-   * @param data - The sibling path data.
+   * @param pathSize - The size of the sibling path.
+   * @param path - The sibling path data.
    */
   constructor(
+    /**
+     * Size of the sibling path (number of fields it contains).
+     */
+    public pathSize: N,
     /**
      * The sibling path data.
      */
-    public data: Buffer[] = [],
-  ) {}
+    path: Buffer[],
+  ) {
+    this.data = assertLength(path, pathSize);
+  }
 
   /**
    * Serializes this SiblingPath object to a buffer.
@@ -48,13 +63,29 @@ export class SiblingPath {
     return serializeBufferArrayToVector(this.data);
   }
 
+  /**
+   * Returns the path buffer underlying the sibling path.
+   * @returns The Buffer array representation of this object.
+   */
+  public toBufferArray(): Buffer[] {
+    return this.data;
+  }
+
+  /**
+   * Convert the Sibling Path object into an array of field elements.
+   * @returns The field array representation of this object.
+   */
+  public toFieldArray(): Fr[] {
+    return this.data.map(buf => Fr.fromBuffer(buf));
+  }
+
   /**
    * Deserializes a SiblingPath from a buffer.
    * @param buf - A buffer containing the buffer representation of SiblingPath.
    * @param offset - An offset to start deserializing from.
    * @returns A SiblingPath object.
    */
-  static fromBuffer(buf: Buffer, offset = 0): SiblingPath {
+  static fromBuffer(buf: Buffer, offset = 0): SiblingPath<number> {
     const { elem } = SiblingPath.deserialize(buf, offset);
     return elem;
   }
@@ -71,7 +102,8 @@ export class SiblingPath {
       adv: 32,
     });
     const { elem, adv } = deserializeArrayFromVector(deserializePath, buf, offset);
-    return { elem: new SiblingPath(elem), adv };
+    const size = elem.length;
+    return { elem: new SiblingPath(size, elem), adv };
   }
 
   /**
@@ -87,7 +119,21 @@ export class SiblingPath {
    * @param repr - A hex string representation of the sibling path.
    * @returns A SiblingPath object.
    */
-  public static fromString(repr: string): SiblingPath {
+  public static fromString(repr: string): SiblingPath<number> {
     return SiblingPath.fromBuffer(Buffer.from(repr, 'hex'));
   }
+
+  /**
+   * Generate a subtree path from the current sibling path.
+   * @param subtreeHeight - The size of the subtree that we are getting the path for.
+   * @returns A new sibling path that is the for the requested subtree.
+   */
+  public getSubtreeSiblingPath<SubtreeHeight extends number, SubtreeSiblingPathHeight extends number>(
+    subtreeHeight: SubtreeHeight,
+  ): SiblingPath<SubtreeSiblingPathHeight> {
+    // Drop the size of the subtree from the path, and return the rest.
+    const subtreeData = this.data.slice(subtreeHeight);
+    const subtreePathSize = (this.pathSize - subtreeHeight) as SubtreeSiblingPathHeight;
+    return new SiblingPath(subtreePathSize, subtreeData);
+  }
 }

yarn-project/merkle-tree/src/sparse_tree/sparse_tree.test.ts

@@ -28,6 +28,8 @@ const createFromName = async (levelUp: levelup.LevelUp, hasher: Hasher, name: st
   return await loadTree(SparseTree, levelUp, hasher, name);
 };
 
+const TEST_TREE_DEPTH = 3;
+
 treeTestSuite('SparseTree', createDb, createFromName);
 standardBasedTreeTestSuite('SparseTree', createDb);
 
@@ -108,7 +110,7 @@ describe('SparseTreeSpecific', () => {
       const root = pedersen.compress(level1LeftHash, level1ZeroHash);
       expect(tree.getRoot(true)).toEqual(root);
       expect(await tree.getSiblingPath(3n, true)).toEqual(
-        new SiblingPath([INITIAL_LEAF, level2ZeroHash, level1ZeroHash]),
+        new SiblingPath(TEST_TREE_DEPTH, [INITIAL_LEAF, level2ZeroHash, level1ZeroHash]),
       );
     }
 
@@ -123,7 +125,7 @@ describe('SparseTreeSpecific', () => {
       const root = pedersen.compress(level1LeftHash, level1RightHash);
       expect(tree.getRoot(true)).toEqual(root);
       expect(await tree.getSiblingPath(6n, true)).toEqual(
-        new SiblingPath([INITIAL_LEAF, level2ZeroHash, level1LeftHash]),
+        new SiblingPath(TEST_TREE_DEPTH, [INITIAL_LEAF, level2ZeroHash, level1LeftHash]),
       );
     }
 
@@ -137,7 +139,7 @@ describe('SparseTreeSpecific', () => {
       const root = pedersen.compress(level1LeftHash, level1RightHash);
       expect(tree.getRoot(true)).toEqual(root);
       expect(await tree.getSiblingPath(2n, true)).toEqual(
-        new SiblingPath([leafAtIndex3, level2ZeroHash, level1RightHash]),
+        new SiblingPath(TEST_TREE_DEPTH, [leafAtIndex3, level2ZeroHash, level1RightHash]),
       );
     }
 
@@ -151,7 +153,7 @@ describe('SparseTreeSpecific', () => {
       const root = pedersen.compress(level1LeftHash, level1RightHash);
       expect(tree.getRoot(true)).toEqual(root);
       expect(await tree.getSiblingPath(3n, true)).toEqual(
-        new SiblingPath([leafAtIndex2, level2ZeroHash, level1RightHash]),
+        new SiblingPath(TEST_TREE_DEPTH, [leafAtIndex2, level2ZeroHash, level1RightHash]),
       );
     }
   });