test: Ephemeral tree hashPreimage

yarn-project/simulator/src/avm/avm_tree.test.ts

@@ -567,6 +567,7 @@ describe('Batch Insertion', () => {
   });
 });
 
+// This benchmark also performs a convenient sanity check
 /* eslint no-console: ["error", { allow: ["time", "timeEnd"] }] */
 describe('A basic benchmark', () => {
   it('Should benchmark writes', async () => {
@@ -576,14 +577,30 @@ describe('A basic benchmark', () => {
     const slots = leaves.map((_, i) => new Fr(i + 128));
 
     const container = await AvmEphemeralForest.create(copyState);
+    await publicDataInsertWorldState(new Fr(0), new Fr(128));
 
     // Updating the first slot, triggers the index 0 to be added to the minimum stored key in the container
     await container.writePublicStorage(new Fr(0), new Fr(128));
+
+    // Check Roots before benchmarking
+    let wsRoot = await getWorldStateRoot(MerkleTreeId.PUBLIC_DATA_TREE);
+    let computedRoot = container.treeMap.get(MerkleTreeId.PUBLIC_DATA_TREE)!.getRoot();
+    expect(computedRoot.toBuffer()).toEqual(wsRoot);
+
     console.time('benchmark');
     // These writes are all new leaves and should be impacted by the key sorted algorithm of the tree.
     for (let i = 0; i < leaves.length; i++) {
       await container.writePublicStorage(slots[i], leaves[i]);
     }
     console.timeEnd('benchmark');
+
+    // Update worldstate for sanity check
+    for (let i = 0; i < leaves.length; i++) {
+      await publicDataInsertWorldState(slots[i], leaves[i]);
+    }
+    // Check roots
+    wsRoot = await getWorldStateRoot(MerkleTreeId.PUBLIC_DATA_TREE);
+    computedRoot = container.treeMap.get(MerkleTreeId.PUBLIC_DATA_TREE)!.getRoot();
+    expect(computedRoot.toBuffer()).toEqual(wsRoot);
   });
 });

yarn-project/simulator/src/avm/avm_tree.ts

@@ -408,15 +408,15 @@ export class AvmEphemeralForest {
         // We are starting with the leaf with largest key <= the specified key
         // Starting at that "min leaf", search for specified key in both the indexed updates
         // and the underlying DB. If not found, return its low leaf.
-        const leafOrLowLeafInfo = await this._searchForLeafOrLowLeaf<ID, T>(
+        const [leafOrLowLeafInfo, pathAbsentInEphemeralTree] = await this._searchForLeafOrLowLeaf<ID, T>(
           treeId,
           bigIntKey,
           minPreimage,
           minIndexedLeafIndex,
         );
         // We did not find it - this is unexpected... the leaf OR low leaf should always be present
         assert(leafOrLowLeafInfo !== undefined, 'Could not find leaf or low leaf. This should not happen!');
-        return [leafOrLowLeafInfo, /*pathAbsentInEphemeralTree=*/ false];
+        return [leafOrLowLeafInfo, pathAbsentInEphemeralTree];
       }
     }
   }
@@ -480,6 +480,10 @@ export class AvmEphemeralForest {
    * @param minPreimage - The leaf with the largest key <= the specified key. Expected to be present in local indexedUpdates.
    * @param minIndex - The index of the leaf with the largest key <= the specified key.
    * @param T - The type of the preimage (PublicData or Nullifier)
+   * @returns [
+   *     preimageWitness | undefined - The leaf or low leaf info (preimage & leaf index),
+   *     pathAbsentInEphemeralTree - whether its sibling path is absent in the ephemeral tree (useful during insertions)
+   * ]
    *
    * @details We look for the low element by bouncing between our local indexedUpdates map or the external DB
    * The conditions we are looking for are:
@@ -493,14 +497,15 @@ export class AvmEphemeralForest {
     key: bigint,
     minPreimage: T,
     minIndex: bigint,
-  ): Promise<PreimageWitness<T> | undefined> {
+  ): Promise<[PreimageWitness<T> | undefined, /*pathAbsentInEphemeralTree=*/ boolean]> {
     let found = false;
     let curr = minPreimage as T;
     let result: PreimageWitness<T> | undefined = undefined;
     // Temp to avoid infinite loops - the limit is the number of leaves we may have to read
     const LIMIT = 2n ** BigInt(getTreeHeight(treeId)) - 1n;
     let counter = 0n;
     let lowPublicDataIndex = minIndex;
+    let pathAbsentInEphemeralTree = false;
     while (!found && counter < LIMIT) {
       const bigIntKey = key;
       if (curr.getKey() === bigIntKey) {
@@ -517,27 +522,24 @@ export class AvmEphemeralForest {
         lowPublicDataIndex = curr.getNextIndex();
         if (this.hasLocalUpdates(treeId, lowPublicDataIndex)) {
           curr = this.getIndexedUpdate(treeId, lowPublicDataIndex)!;
+          pathAbsentInEphemeralTree = false;
         } else {
           const preimage: IndexedTreeLeafPreimage = (await this.treeDb.getLeafPreimage(treeId, lowPublicDataIndex))!;
           curr = preimage as T;
+          pathAbsentInEphemeralTree = true;
         }
       }
       counter++;
     }
-    return result;
+    return [result, pathAbsentInEphemeralTree];;
   }
 
   /**
    * This hashes the preimage to a field element
    */
   hashPreimage<T extends TreeLeafPreimage>(preimage: T): Fr {
-    // Watch for this edge-case, we are hashing the key=0 leaf to 0.
-    // This is for backward compatibility with the world state implementation
-    if (preimage.getKey() === 0n) {
-      return Fr.zero();
-    }
-    const input = preimage.toHashInputs().map(x => Fr.fromBuffer(x));
-    return poseidon2Hash(input);
+     const input = preimage.toHashInputs().map(x => Fr.fromBuffer(x));
+     return poseidon2Hash(input);
   }
 }