feat: rework some storage read collection, change update_check, bring in changes from pi PR

Author: MirandaWood

barretenberg/cpp/src/barretenberg/crypto/merkle_tree/response.hpp

@@ -71,6 +71,7 @@ template <typename LeafType> struct LeafUpdateWitnessData {
     LeafUpdateWitnessData(LeafUpdateWitnessData&& other) noexcept = default;
     LeafUpdateWitnessData& operator=(const LeafUpdateWitnessData& other) = default;
     LeafUpdateWitnessData& operator=(LeafUpdateWitnessData&& other) noexcept = default;
+    bool operator==(const LeafUpdateWitnessData& other) const = default;
 
     MSGPACK_FIELDS(leaf, index, path);
 };

barretenberg/cpp/src/barretenberg/vm2/avm_sim_api.cpp

@@ -14,11 +14,6 @@ TxSimulationResult AvmSimAPI::simulate_with_hinted_dbs(const ProvingInputs& inpu
     AvmSimulationHelper simulation_helper;
     auto result = AVM_TRACK_TIME_V("simulation/all", simulation_helper.simulate_fast_with_hinted_dbs(inputs.hints));
 
-    if (debug_logging) {
-        // TODO(fcarreiro): Enable once PI generation is complete.
-        // BB_ASSERT_EQ(inputs.publicInputs, result.public_inputs);
-    }
-
     return result;
 }
 

barretenberg/cpp/src/barretenberg/vm2/simulation/gadgets/tx_execution.hpp

@@ -51,6 +51,8 @@ class TxExecution final {
 
     const TxContext& get_tx_context() const { return tx_context; }
 
+    const TxContext& get_tx_context() const { return tx_context; }
+
   private:
     ExecutionInterface& call_execution;
     ContextProviderInterface& context_provider;

barretenberg/cpp/src/barretenberg/vm2/simulation/gadgets/update_check.cpp

@@ -11,10 +11,12 @@ namespace {
 
 using UnconstrainedPoseidon2 = crypto::Poseidon2<crypto::Poseidon2Bn254ScalarFieldParams>;
 
-FF unconstrained_read(LowLevelMerkleDBInterface& merkle_db, const FF& leaf_slot)
+FF unconstrained_read(const LowLevelMerkleDBInterface& merkle_db, const FF& leaf_slot)
 {
     auto [present, index] = merkle_db.get_low_indexed_leaf(world_state::MerkleTreeId::PUBLIC_DATA_TREE, leaf_slot);
+    info("performing unconstrained read for slot: ", leaf_slot, ", w res index: ", index);
     auto preimage = merkle_db.get_leaf_preimage_public_data_tree(index);
+    info("performing get leaf preimage for index: ", index, ", got value: ", preimage.leaf.value);
     return present ? preimage.leaf.value : 0;
 }
 
@@ -23,6 +25,7 @@ FF unconstrained_read(LowLevelMerkleDBInterface& merkle_db, const FF& leaf_slot)
 void UpdateCheck::check_current_class_id(const AztecAddress& address, const ContractInstance& instance)
 {
     // Compute the public data tree slots
+    info("should arrive here after instance");
     FF delayed_public_mutable_slot = poseidon2.hash({ UPDATED_CLASS_IDS_SLOT, address });
     FF delayed_public_mutable_hash_slot = delayed_public_mutable_slot + UPDATES_DELAYED_PUBLIC_MUTABLE_VALUES_LEN;
     // Read the hash from the tree. We do a trick with delayed public mutables (updates are delayed public mutables)
@@ -38,6 +41,13 @@ void UpdateCheck::check_current_class_id(const AztecAddress& address, const Cont
     uint64_t current_timestamp = globals.timestamp;
 
     if (hash == 0) {
+        info("hash 0, not performing unconstrained reads, would be for slots:");
+        info(unconstrained_compute_leaf_slot(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS, delayed_public_mutable_slot));
+        info(unconstrained_compute_leaf_slot(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
+                                             delayed_public_mutable_slot + 1));
+        info(unconstrained_compute_leaf_slot(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
+                                             delayed_public_mutable_slot + 2));
+
         // If the delayed public mutable has never been written, then the contract was never updated. We short circuit
         // early.
         if (instance.original_class_id != instance.current_class_id) {
@@ -52,6 +62,7 @@ void UpdateCheck::check_current_class_id(const AztecAddress& address, const Cont
         for (size_t i = 0; i < update_preimage.size(); ++i) {
             FF leaf_slot = unconstrained_compute_leaf_slot(CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
                                                            delayed_public_mutable_slot + i);
+            info("performing unconstrained reads i = ", i);
             update_preimage[i] = unconstrained_read(unconstrained_merkle_db, leaf_slot);
         }
 

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

@@ -7,7 +7,6 @@
 #include <string>
 #include <vector>
 
-#include "barretenberg/common/bb_bench.hpp"
 #include "barretenberg/common/log.hpp"
 #include "barretenberg/vm2/common/avm_inputs.hpp"
 #include "barretenberg/vm2/common/aztec_types.hpp"
@@ -127,6 +126,22 @@ GetLowIndexedLeafResponse HintingRawDB::get_low_indexed_leaf(world_state::Merkle
     query_hints.get_previous_value_index_hints[key] = { resp.is_already_present, resp.index };
     // TODO(MW): We may need a sibling path hint so must collect it in case - see comments in public_db_sources.ts
     get_sibling_path(tree_id, resp.index);
+
+    if (tree_id == world_state::MerkleTreeId::NULLIFIER_TREE) {
+        // TODO(MW): We may need a GetLeafPreimageHint for the nullifier tree when calling nullifier_exists, so collect
+        // it
+        // in case.
+        // NB: The PureMerkleDB does not perform this, but the nullifier check gadget requires a leaf preimage. Ts
+        // gathers the hint: (state_manager -> checkNullifierExists() -> doMerkleOperations -> public_db_sources ->
+        // checkNullifierExists())
+        get_leaf_preimage_nullifier_tree(resp.index);
+    } else if ((tree_id == world_state::MerkleTreeId::PUBLIC_DATA_TREE) && (!resp.is_already_present)) {
+        // TODO(MW): We may need a GetLeafPreimageHint for the public data tree when calling storage_read, so collect it
+        // in case.
+        // NB: The PureMerkleDB does not perform this if !is_already_present, but MerkleDB and ts perform it
+        // unconditionally. Ts gathers the hint: (public_db_sources -> storageRead())
+        get_leaf_preimage_public_data_tree(resp.index);
+    }
     return resp;
 }
 
@@ -286,14 +301,14 @@ std::vector<AppendLeafResult> HintingRawDB::append_leaves(world_state::MerkleTre
     for (uint32_t i = 0; i < leaves.size(); i++) {
         FF root_after = i == leaves.size() - 1 ? get_tree_info(tree_id).root : results[i + 1].root;
         // Iterate tree_info to the be state after adding this leaf:
-        tree_info = appendLeafInternal(tree_info, root_after, tree_id, leaves[i]);
+        tree_info = appendLeafInternal(tree_info, results[i].path, root_after, tree_id, leaves[i]);
     }
 
     return results;
 }
 
-// TODO(MW): rework
 AppendOnlyTreeSnapshot HintingRawDB::appendLeafInternal(AppendOnlyTreeSnapshot state_before,
+                                                        SiblingPath& path,
                                                         const FF& root_after,
                                                         world_state::MerkleTreeId tree_id,
                                                         const FF& leaf)
@@ -303,14 +318,15 @@ AppendOnlyTreeSnapshot HintingRawDB::appendLeafInternal(AppendOnlyTreeSnapshot s
     // constraint the insertion.
     // TODO(https://github.com/AztecProtocol/aztec-packages/issues/13380): This can be changed if the world state
     // appendLeaves returns the sibling paths.
-    AppendLeavesHintKey key = { state_before, tree_id, { leaf } };
+    AppendLeavesHintKey append_key = { state_before, tree_id, { leaf } };
     AppendOnlyTreeSnapshot state_after = { .root = root_after,
                                            .nextAvailableLeafIndex = state_before.nextAvailableLeafIndex + 1 };
-    append_leaves_hints[key] = state_after;
-    // TODO(MW): Just store hint using result in append_leaves to avoid unnecessary extra calls to underlying db?
-    // NOTE: this call will use the /current/ db.get_tree_info() (post full append_leaves), which may not match that at
-    // root_after:
-    get_sibling_path(tree_id, state_before.nextAvailableLeafIndex);
+    append_leaves_hints[append_key] = state_after;
+    // TODO(MW): Storing sibling path hint manually using the result since a get_sibling_path() call here will use the
+    // /current/ db.get_tree_info() (post full append_leaves), which may not match that at result.root. We may not care
+    // about this (see comment in PureRawMerkleDB::append_leaves())
+    GetSiblingPathKey path_key = { state_after, tree_id, state_before.nextAvailableLeafIndex };
+    query_hints.get_sibling_path_hints[path_key] = path;
     return state_after;
 }
 

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

@@ -101,6 +101,7 @@ class HintingRawDB final : public LowLevelMerkleDBInterface {
     // Private helper methods. TODO(MW): extract out? Copied from raw_data_dbs
     const AppendOnlyTreeSnapshot& get_tree_info(MerkleTreeId tree_id) const;
     AppendOnlyTreeSnapshot appendLeafInternal(AppendOnlyTreeSnapshot state_before,
+                                              SiblingPath& path,
                                               const FF& root_after,
                                               MerkleTreeId tree_id,
                                               const FF& leaf);