feat(avm): merkle hints (part 2)

barretenberg/cpp/src/barretenberg/vm2/common/avm_inputs.hpp

@@ -147,6 +147,19 @@ template <typename Leaf> struct GetLeafPreimageHint {
     MSGPACK_FIELDS(hintKey, index, leaf, nextIndex, nextValue);
 };
 
+struct GetLeafValueHint {
+    AppendOnlyTreeSnapshot hintKey;
+    // params
+    world_state::MerkleTreeId treeId;
+    uint64_t index;
+    // return
+    FF value;
+
+    bool operator==(const GetLeafValueHint& other) const = default;
+
+    MSGPACK_FIELDS(hintKey, treeId, index, value);
+};
+
 ////////////////////////////////////////////////////////////////////////////
 // Hints (other)
 ////////////////////////////////////////////////////////////////////////////
@@ -174,6 +187,8 @@ struct ExecutionHints {
     std::vector<GetSiblingPathHint> getSiblingPathHints;
     std::vector<GetPreviousValueIndexHint> getPreviousValueIndexHints;
     std::vector<GetLeafPreimageHint<crypto::merkle_tree::PublicDataLeafValue>> getLeafPreimageHintsPublicDataTree;
+    std::vector<GetLeafPreimageHint<crypto::merkle_tree::NullifierLeafValue>> getLeafPreimageHintsNullifierTree;
+    std::vector<GetLeafValueHint> getLeafValueHints;
 
     bool operator==(const ExecutionHints& other) const = default;
 
@@ -183,7 +198,9 @@ struct ExecutionHints {
                    bytecodeCommitments,
                    getSiblingPathHints,
                    getPreviousValueIndexHints,
-                   getLeafPreimageHintsPublicDataTree);
+                   getLeafPreimageHintsPublicDataTree,
+                   getLeafPreimageHintsNullifierTree,
+                   getLeafValueHints);
 };
 
 ////////////////////////////////////////////////////////////////////////////

barretenberg/cpp/src/barretenberg/vm2/simulation/lib/db_interfaces.hpp

@@ -3,6 +3,7 @@
 #include <optional>
 
 #include "barretenberg/crypto/merkle_tree/hash_path.hpp"
+#include "barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp"
 #include "barretenberg/crypto/merkle_tree/response.hpp"
 #include "barretenberg/crypto/merkle_tree/types.hpp"
 #include "barretenberg/vm2/common/avm_inputs.hpp"
@@ -30,9 +31,13 @@ class LowLevelMerkleDBInterface {
                                                                   crypto::merkle_tree::index_t leaf_index) const = 0;
     virtual crypto::merkle_tree::GetLowIndexedLeafResponse get_low_indexed_leaf(world_state::MerkleTreeId tree_id,
                                                                                 const FF& value) const = 0;
+    // Returns the value if it exists, 0 otherwise.
+    virtual FF get_leaf_value(world_state::MerkleTreeId tree_id, crypto::merkle_tree::index_t leaf_index) const = 0;
     // We don't template the preimage methods because templated methods cannot be virtual.
     virtual crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue>
     get_leaf_preimage_public_data_tree(crypto::merkle_tree::index_t leaf_index) const = 0;
+    virtual crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::NullifierLeafValue> get_leaf_preimage_nullifier_tree(
+        crypto::merkle_tree::index_t leaf_index) const = 0;
 };
 
 // High level access to a merkle db. In general these will be constrained.

barretenberg/cpp/src/barretenberg/vm2/simulation/lib/raw_data_dbs.cpp

@@ -5,6 +5,7 @@
 #include <stdexcept>
 
 #include "barretenberg/common/log.hpp"
+#include "barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp"
 #include "barretenberg/vm2/simulation/lib/contract_crypto.hpp"
 
 namespace bb::avm2::simulation {
@@ -101,7 +102,11 @@ HintedRawMerkleDB::HintedRawMerkleDB(const ExecutionHints& hints, const TreeSnap
           "\n * get_previous_value_index hints: ",
           hints.getPreviousValueIndexHints.size(),
           "\n * get_leaf_preimage hints_public_data_tree: ",
-          hints.getLeafPreimageHintsPublicDataTree.size());
+          hints.getLeafPreimageHintsPublicDataTree.size(),
+          "\n * get_leaf_preimage hints_nullifier_tree: ",
+          hints.getLeafPreimageHintsNullifierTree.size(),
+          "\n * get_leaf_value_hints: ",
+          hints.getLeafValueHints.size());
     debug("Initializing HintedRawMerkleDB with snapshots...",
           "\n * nullifierTree: ",
           tree_roots.nullifierTree.root,
@@ -149,6 +154,20 @@ HintedRawMerkleDB::HintedRawMerkleDB(const ExecutionHints& hints, const TreeSnap
             /*nextVal=*/get_leaf_preimage_hint.nextValue,
         };
     }
+
+    for (const auto& get_leaf_preimage_hint : hints.getLeafPreimageHintsNullifierTree) {
+        GetLeafPreimageKey key = { get_leaf_preimage_hint.hintKey, get_leaf_preimage_hint.index };
+        get_leaf_preimage_hints_nullifier_tree[key] = {
+            /*val=*/get_leaf_preimage_hint.leaf,
+            /*nextIdx=*/get_leaf_preimage_hint.nextIndex,
+            /*nextVal=*/get_leaf_preimage_hint.nextValue,
+        };
+    }
+
+    for (const auto& get_leaf_value_hint : hints.getLeafValueHints) {
+        GetLeafValueKey key = { get_leaf_value_hint.hintKey, get_leaf_value_hint.treeId, get_leaf_value_hint.index };
+        get_leaf_value_hints[key] = get_leaf_value_hint.value;
+    }
 }
 
 const AppendOnlyTreeSnapshot& HintedRawMerkleDB::get_tree_info(world_state::MerkleTreeId tree_id) const
@@ -207,6 +226,14 @@ crypto::merkle_tree::GetLowIndexedLeafResponse HintedRawMerkleDB::get_low_indexe
     return it->second;
 }
 
+FF HintedRawMerkleDB::get_leaf_value(world_state::MerkleTreeId tree_id, crypto::merkle_tree::index_t leaf_index) const
+{
+    auto tree_info = get_tree_info(tree_id);
+    GetLeafValueKey key = { tree_info, tree_id, leaf_index };
+    auto it = get_leaf_value_hints.find(key);
+    return it == get_leaf_value_hints.end() ? 0 : it->second;
+}
+
 crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue> HintedRawMerkleDB::
     get_leaf_preimage_public_data_tree(crypto::merkle_tree::index_t leaf_index) const
 {
@@ -225,4 +252,22 @@ crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue> Hinte
     return it->second;
 }
 
+crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::NullifierLeafValue> HintedRawMerkleDB::
+    get_leaf_preimage_nullifier_tree(crypto::merkle_tree::index_t leaf_index) const
+{
+    auto tree_info = get_tree_info(world_state::MerkleTreeId::NULLIFIER_TREE);
+    GetLeafPreimageKey key = { tree_info, leaf_index };
+    auto it = get_leaf_preimage_hints_nullifier_tree.find(key);
+    if (it == get_leaf_preimage_hints_nullifier_tree.end()) {
+        throw std::runtime_error(format("Leaf preimage (NULLIFIER_TREE) not found for key (root: ",
+                                        tree_info.root,
+                                        ", size: ",
+                                        tree_info.nextAvailableLeafIndex,
+                                        ", leaf_index: ",
+                                        leaf_index,
+                                        ")"));
+    }
+    return it->second;
+}
+
 } // namespace bb::avm2::simulation

barretenberg/cpp/src/barretenberg/vm2/simulation/lib/raw_data_dbs.hpp

@@ -2,6 +2,7 @@
 
 #include "barretenberg/common/utils.hpp"
 #include "barretenberg/crypto/merkle_tree/hash_path.hpp"
+#include "barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp"
 #include "barretenberg/vm2/common/avm_inputs.hpp"
 #include "barretenberg/vm2/common/aztec_types.hpp"
 #include "barretenberg/vm2/common/field.hpp"
@@ -40,8 +41,11 @@ class HintedRawMerkleDB final : public LowLevelMerkleDBInterface {
                                                           crypto::merkle_tree::index_t leaf_index) const override;
     crypto::merkle_tree::GetLowIndexedLeafResponse get_low_indexed_leaf(world_state::MerkleTreeId tree_id,
                                                                         const FF& value) const override;
+    FF get_leaf_value(world_state::MerkleTreeId tree_id, crypto::merkle_tree::index_t leaf_index) const override;
     crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue> get_leaf_preimage_public_data_tree(
         crypto::merkle_tree::index_t leaf_index) const override;
+    crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::NullifierLeafValue> get_leaf_preimage_nullifier_tree(
+        crypto::merkle_tree::index_t leaf_index) const override;
 
   private:
     TreeSnapshots tree_roots;
@@ -55,6 +59,11 @@ class HintedRawMerkleDB final : public LowLevelMerkleDBInterface {
     using GetLeafPreimageKey = utils::HashableTuple<AppendOnlyTreeSnapshot, crypto::merkle_tree::index_t>;
     unordered_flat_map<GetLeafPreimageKey, crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue>>
         get_leaf_preimage_hints_public_data_tree;
+    unordered_flat_map<GetLeafPreimageKey, crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::NullifierLeafValue>>
+        get_leaf_preimage_hints_nullifier_tree;
+    using GetLeafValueKey =
+        utils::HashableTuple<AppendOnlyTreeSnapshot, world_state::MerkleTreeId, crypto::merkle_tree::index_t>;
+    unordered_flat_map<GetLeafValueKey, FF> get_leaf_value_hints;
 
     const AppendOnlyTreeSnapshot& get_tree_info(world_state::MerkleTreeId tree_id) const;
 };

barretenberg/cpp/src/barretenberg/vm2/simulation/testing/mock_dbs.hpp

@@ -32,10 +32,18 @@ class MockLowLevelMerkleDB : public LowLevelMerkleDBInterface {
                 get_low_indexed_leaf,
                 (world_state::MerkleTreeId tree_id, const FF& value),
                 (const, override));
+    MOCK_METHOD(FF,
+                get_leaf_value,
+                (world_state::MerkleTreeId tree_id, crypto::merkle_tree::index_t leaf_index),
+                (const, override));
     MOCK_METHOD(crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::PublicDataLeafValue>,
                 get_leaf_preimage_public_data_tree,
                 (crypto::merkle_tree::index_t leaf_index),
                 (const, override));
+    MOCK_METHOD(crypto::merkle_tree::IndexedLeaf<crypto::merkle_tree::NullifierLeafValue>,
+                get_leaf_preimage_nullifier_tree,
+                (crypto::merkle_tree::index_t leaf_index),
+                (const, override));
 };
 
 class MockHighLevelMerkleDB : public HighLevelMerkleDBInterface {

yarn-project/simulator/src/public/hinting_db_sources.ts

@@ -7,7 +7,9 @@ import {
   AvmContractClassHint,
   AvmContractInstanceHint,
   type AvmExecutionHints,
+  AvmGetLeafPreimageHintNullifierTree,
   AvmGetLeafPreimageHintPublicDataTree,
+  AvmGetLeafValueHint,
   AvmGetPreviousValueIndexHint,
   AvmGetSiblingPathHint,
 } from '@aztec/stdlib/avm';
@@ -17,11 +19,15 @@ import {
   AppendOnlyTreeSnapshot,
   type IndexedTreeId,
   MerkleTreeId,
+  type MerkleTreeLeafType,
+  NullifierLeaf,
   PublicDataTreeLeaf,
   type SequentialInsertionResult,
   getTreeName,
 } from '@aztec/stdlib/trees';
 
+import { strict as assert } from 'assert';
+
 import type { PublicContractsDBInterface } from '../common/db_interfaces.js';
 import { PublicTreesDB } from './public_db_sources.js';
 
@@ -147,15 +153,44 @@ export class HintingPublicTreesDB extends PublicTreesDB {
             ),
           );
           break;
+        case MerkleTreeId.NULLIFIER_TREE:
+          this.hints.getLeafPreimageHintsNullifierTree.push(
+            new AvmGetLeafPreimageHintNullifierTree(
+              key,
+              index,
+              preimage.asLeaf() as NullifierLeaf,
+              preimage.getNextIndex(),
+              new Fr(preimage.getNextKey()),
+            ),
+          );
+          break;
         default:
-          HintingPublicTreesDB.log.debug(`getLeafPreimage not hinted for tree ${getTreeName(treeId)} yet!`);
+          // Use getLeafValue for the other trees.
+          throw new Error('getLeafPreimage only supported for PublicDataTree and NullifierTree!');
           break;
       }
     }
 
     return preimage;
   }
 
+  public override async getLeafValue<ID extends MerkleTreeId>(
+    treeId: ID,
+    index: bigint,
+  ): Promise<MerkleTreeLeafType<typeof treeId> | undefined> {
+    // Use getLeafPreimage for PublicDataTree and NullifierTree.
+    assert(treeId == MerkleTreeId.NOTE_HASH_TREE || treeId == MerkleTreeId.L1_TO_L2_MESSAGE_TREE);
+
+    const value = await super.getLeafValue<ID>(treeId, index);
+    if (value) {
+      const key = await this.getHintKey(treeId);
+      // We can cast to Fr because we know the type of the tree.
+      this.hints.getLeafValueHints.push(new AvmGetLeafValueHint(key, treeId, index, value as Fr));
+    }
+
+    return value;
+  }
+
   // State modification.
   public override async sequentialInsert<TreeHeight extends number, ID extends IndexedTreeId>(
     treeId: ID,

yarn-project/simulator/src/public/public_db_sources.ts

@@ -437,6 +437,9 @@ export class PublicTreesDB extends ForwardMerkleTree implements PublicStateDBInt
   public async getL1ToL2LeafValue(leafIndex: bigint): Promise<Fr | undefined> {
     const timer = new Timer();
     const leafValue = await this.getLeafValue(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, leafIndex);
+    // TODO(fcarreiro): We need this for the hints. Might move it to the hinting layer.
+    await this.getSiblingPath(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, leafIndex);
+
     this.logger.debug(`[DB] Fetched L1 to L2 message leaf value`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
@@ -448,6 +451,9 @@ export class PublicTreesDB extends ForwardMerkleTree implements PublicStateDBInt
   public async getNoteHash(leafIndex: bigint): Promise<Fr | undefined> {
     const timer = new Timer();
     const leafValue = await this.getLeafValue(MerkleTreeId.NOTE_HASH_TREE, leafIndex);
+    // TODO(fcarreiro): We need this for the hints. Might move it to the hinting layer.
+    await this.getSiblingPath(MerkleTreeId.NOTE_HASH_TREE, leafIndex);
+
     this.logger.debug(`[DB] Fetched note hash leaf value`, {
       eventName: 'public-db-access',
       duration: timer.ms(),
@@ -458,7 +464,16 @@ export class PublicTreesDB extends ForwardMerkleTree implements PublicStateDBInt
 
   public async getNullifierIndex(nullifier: Fr): Promise<bigint | undefined> {
     const timer = new Timer();
-    const index = (await this.findLeafIndices(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()]))[0];
+    const lowLeafResult = await this.getPreviousValueIndex(MerkleTreeId.NULLIFIER_TREE, nullifier.toBigInt());
+    if (!lowLeafResult) {
+      throw new Error('Low leaf not found');
+    }
+    // TODO(fcarreiro): We need this for the hints. Might move it to the hinting layer.
+    await this.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, lowLeafResult.index);
+    // TODO(fcarreiro): We need this for the hints. Might move it to the hinting layer.
+    await this.getLeafPreimage(MerkleTreeId.NULLIFIER_TREE, lowLeafResult.index);
+    const index = lowLeafResult.alreadyPresent ? lowLeafResult.index : undefined;
+
     this.logger.debug(`[DB] Fetched nullifier index`, {
       eventName: 'public-db-access',
       duration: timer.ms(),

yarn-project/stdlib/src/avm/avm.ts

@@ -9,6 +9,7 @@ import { AztecAddress } from '../aztec-address/index.js';
 import { PublicKeys } from '../keys/public_keys.js';
 import { AppendOnlyTreeSnapshot } from '../trees/append_only_tree_snapshot.js';
 import type { MerkleTreeId } from '../trees/merkle_tree_id.js';
+import { NullifierLeaf } from '../trees/nullifier_leaf.js';
 import { PublicDataTreeLeaf } from '../trees/public_data_leaf.js';
 import { AvmCircuitPublicInputs } from './avm_circuit_public_inputs.js';
 import { serializeWithMessagePack } from './message_pack.js';
@@ -186,7 +187,33 @@ function AvmGetLeafPreimageHintFactory<T extends IndexedTreeLeaf>(klass: {
   };
 }
 
+// Note: only supported for PUBLIC_DATA_TREE and NULLIFIER_TREE.
 export class AvmGetLeafPreimageHintPublicDataTree extends AvmGetLeafPreimageHintFactory(PublicDataTreeLeaf) {}
+export class AvmGetLeafPreimageHintNullifierTree extends AvmGetLeafPreimageHintFactory(NullifierLeaf) {}
+
+// Hint for MerkleTreeDB.getLeafValue.
+// Note: only supported for NOTE_HASH_TREE and L1_TO_L2_MESSAGE_TREE.
+export class AvmGetLeafValueHint {
+  constructor(
+    public readonly hintKey: AppendOnlyTreeSnapshot,
+    // params
+    public readonly treeId: MerkleTreeId,
+    public readonly index: bigint,
+    // return
+    public readonly value: Fr,
+  ) {}
+
+  static get schema() {
+    return z
+      .object({
+        hintKey: AppendOnlyTreeSnapshot.schema,
+        treeId: z.number().int().nonnegative(),
+        index: schemas.BigInt,
+        value: schemas.Fr,
+      })
+      .transform(({ hintKey, treeId, index, value }) => new AvmGetLeafValueHint(hintKey, treeId, index, value));
+  }
+}
 
 ////////////////////////////////////////////////////////////////////////////
 // Hints (other)
@@ -225,6 +252,8 @@ export class AvmExecutionHints {
     public readonly getSiblingPathHints: AvmGetSiblingPathHint[] = [],
     public readonly getPreviousValueIndexHints: AvmGetPreviousValueIndexHint[] = [],
     public readonly getLeafPreimageHintsPublicDataTree: AvmGetLeafPreimageHintPublicDataTree[] = [],
+    public readonly getLeafPreimageHintsNullifierTree: AvmGetLeafPreimageHintNullifierTree[] = [],
+    public readonly getLeafValueHints: AvmGetLeafValueHint[] = [],
   ) {}
 
   static empty() {
@@ -241,6 +270,8 @@ export class AvmExecutionHints {
         getSiblingPathHints: AvmGetSiblingPathHint.schema.array(),
         getPreviousValueIndexHints: AvmGetPreviousValueIndexHint.schema.array(),
         getLeafPreimageHintsPublicDataTree: AvmGetLeafPreimageHintPublicDataTree.schema.array(),
+        getLeafPreimageHintsNullifierTree: AvmGetLeafPreimageHintNullifierTree.schema.array(),
+        getLeafValueHints: AvmGetLeafValueHint.schema.array(),
       })
       .transform(
         ({
@@ -251,6 +282,8 @@ export class AvmExecutionHints {
           getSiblingPathHints,
           getPreviousValueIndexHints,
           getLeafPreimageHintsPublicDataTree,
+          getLeafPreimageHintsNullifierTree,
+          getLeafValueHints,
         }) =>
           new AvmExecutionHints(
             enqueuedCalls,
@@ -260,6 +293,8 @@ export class AvmExecutionHints {
             getSiblingPathHints,
             getPreviousValueIndexHints,
             getLeafPreimageHintsPublicDataTree,
+            getLeafPreimageHintsNullifierTree,
+            getLeafValueHints,
           ),
       );
   }