From 887c01103255ea4cbbb6cb33c8771d47123b3bff Mon Sep 17 00:00:00 2001 From: ludamad Date: Fri, 29 Nov 2024 14:49:44 -0500 Subject: [PATCH] fix: Revert "feat: Avoid inserting an empty leaf in indexed trees on update" (#10319) Broke kind tests. Reverts AztecProtocol/aztec-packages#10281 --- .../content_addressed_indexed_tree.hpp | 120 +++++++----------- .../merkle_tree/indexed_tree/indexed_leaf.hpp | 2 +- .../vm/avm/trace/gadgets/merkle_tree.cpp | 8 +- .../src/base/components/public_data_tree.nr | 15 ++- .../types/src/data/public_data_tree_leaf.nr | 8 +- .../types/src/merkle_tree/indexed_tree.nr | 52 ++++---- .../indexed_tree/check_valid_low_leaf.nr | 7 - .../types/src/merkle_tree/leaf_preimage.nr | 8 +- .../types/src/merkle_tree/membership.nr | 7 - .../simulator/src/avm/avm_simulator.test.ts | 6 +- .../simulator/src/avm/avm_tree.test.ts | 87 ++++++------- yarn-project/simulator/src/avm/avm_tree.ts | 113 ++++++++--------- .../simulator/src/avm/journal/journal.ts | 25 ++-- 13 files changed, 191 insertions(+), 267 deletions(-) diff --git a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/content_addressed_indexed_tree.hpp b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/content_addressed_indexed_tree.hpp index e1a47ec5351..197f784d59f 100644 --- a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/content_addressed_indexed_tree.hpp +++ b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/content_addressed_indexed_tree.hpp @@ -252,20 +252,18 @@ class ContentAddressedIndexedTree : public ContentAddressedAppendOnlyTree> new_leaf; + IndexedLeafValueType new_leaf; + index_t new_leaf_index; }; struct SequentialInsertionGenerationResponse { - std::vector updates_to_perform; + std::shared_ptr> updates_to_perform; index_t highest_index; }; using SequentialInsertionGenerationCallback = - std::function&)>; + std::function&)>; void generate_sequential_insertions(const std::vector& values, const SequentialInsertionGenerationCallback& completion); @@ -1424,14 +1422,6 @@ void ContentAddressedIndexedTree::add_or_update_values_seq const AddSequentiallyCompletionCallbackWithWitness& completion, bool capture_witness) { - - // This struct is used to collect some state from the asynchronous operations we are about to perform - struct IntermediateResults { - std::vector updates_to_perform; - size_t appended_leaves = 0; - }; - auto results = std::make_shared(); - auto on_error = [=](const std::string& message) { try { TypedResponse> response; @@ -1453,7 +1443,7 @@ void ContentAddressedIndexedTree::add_or_update_values_seq ReadTransactionPtr tx = store_->create_read_transaction(); store_->get_meta(meta, *tx, true); - index_t new_total_size = results->appended_leaves + meta.size; + index_t new_total_size = values.size() + meta.size; meta.size = new_total_size; meta.root = store_->get_current_root(*tx, true); @@ -1462,29 +1452,23 @@ void ContentAddressedIndexedTree::add_or_update_values_seq if (capture_witness) { // Split results->update_witnesses between low_leaf_witness_data and insertion_witness_data + // Currently we always insert an empty leaf, even if it's an update, so we can split based + // on the index of the witness data response.inner.insertion_witness_data = std::make_shared>>(); - response.inner.insertion_witness_data->reserve(results->updates_to_perform.size()); - + ; response.inner.low_leaf_witness_data = std::make_shared>>(); - response.inner.low_leaf_witness_data->reserve(results->updates_to_perform.size()); - - size_t current_witness_index = 0; - for (size_t i = 0; i < results->updates_to_perform.size(); ++i) { - LeafUpdateWitnessData low_leaf_witness = - updates_completion_response.inner.update_witnesses->at(current_witness_index++); - response.inner.low_leaf_witness_data->push_back(low_leaf_witness); - - // If this update has an insertion, append the real witness - if (results->updates_to_perform.at(i).new_leaf.has_value()) { - LeafUpdateWitnessData insertion_witness = - updates_completion_response.inner.update_witnesses->at(current_witness_index++); - response.inner.insertion_witness_data->push_back(insertion_witness); + ; + + for (size_t i = 0; i < updates_completion_response.inner.update_witnesses->size(); ++i) { + LeafUpdateWitnessData& witness_data = + updates_completion_response.inner.update_witnesses->at(i); + // If even, it's a low leaf, if odd, it's an insertion witness + if (i % 2 == 0) { + response.inner.low_leaf_witness_data->push_back(witness_data); } else { - // If it's an update, append an empty witness - response.inner.insertion_witness_data->push_back(LeafUpdateWitnessData{ - .leaf = IndexedLeafValueType::empty(), .index = 0, .path = std::vector(depth_) }); + response.inner.insertion_witness_data->push_back(witness_data); } } } @@ -1496,33 +1480,23 @@ void ContentAddressedIndexedTree::add_or_update_values_seq // This signals the completion of the insertion data generation // Here we'll perform all updates to the tree SequentialInsertionGenerationCallback insertion_generation_completed = - [=, this](TypedResponse& insertion_response) { + [=, this](const TypedResponse& insertion_response) { if (!insertion_response.success) { on_error(insertion_response.message); return; } std::shared_ptr> flat_updates = std::make_shared>(); - flat_updates->reserve(insertion_response.inner.updates_to_perform.size() * 2); - - for (size_t i = 0; i < insertion_response.inner.updates_to_perform.size(); ++i) { - InsertionUpdates& insertion_update = insertion_response.inner.updates_to_perform.at(i); + for (size_t i = 0; i < insertion_response.inner.updates_to_perform->size(); ++i) { + InsertionUpdates& insertion_update = insertion_response.inner.updates_to_perform->at(i); flat_updates->push_back(insertion_update.low_leaf_update); - if (insertion_update.new_leaf.has_value()) { - results->appended_leaves++; - IndexedLeafValueType new_leaf; - index_t new_leaf_index = 0; - std::tie(new_leaf, new_leaf_index) = insertion_update.new_leaf.value(); - flat_updates->push_back(LeafUpdate{ - .leaf_index = new_leaf_index, - .updated_leaf = new_leaf, - .original_leaf = IndexedLeafValueType::empty(), - }); - } + flat_updates->push_back(LeafUpdate{ + .leaf_index = insertion_update.new_leaf_index, + .updated_leaf = insertion_update.new_leaf, + .original_leaf = IndexedLeafValueType::empty(), + }); } - // We won't use anymore updates_to_perform - results->updates_to_perform = std::move(insertion_response.inner.updates_to_perform); - assert(insertion_response.inner.updates_to_perform.size() == 0); + if (capture_witness) { perform_updates(flat_updates->size(), flat_updates, final_completion); return; @@ -1540,12 +1514,27 @@ void ContentAddressedIndexedTree::generate_sequential_inse { execute_and_report( [=, this](TypedResponse& response) { + response.inner.highest_index = 0; + response.inner.updates_to_perform = std::make_shared>(); + TreeMeta meta; ReadTransactionPtr tx = store_->create_read_transaction(); store_->get_meta(meta, *tx, true); RequestContext requestContext; requestContext.includeUncommitted = true; + // Ensure that the tree is not going to be overfilled + index_t new_total_size = values.size() + meta.size; + if (new_total_size > max_size_) { + throw std::runtime_error(format("Unable to insert values into tree ", + meta.name, + " new size: ", + new_total_size, + " max size: ", + max_size_)); + } + // The highest index touched will be the last leaf index, since we append a zero for updates + response.inner.highest_index = new_total_size - 1; requestContext.root = store_->get_current_root(*tx, true); // Fetch the frontier (non empty nodes to the right) of the tree. This will ensure that perform_updates or // perform_updates_without_witness has all the cached nodes it needs to perform the insertions. See comment @@ -1554,15 +1543,12 @@ void ContentAddressedIndexedTree::generate_sequential_inse find_leaf_hash(meta.size - 1, requestContext, *tx, true); } - index_t current_size = meta.size; - for (size_t i = 0; i < values.size(); ++i) { const LeafValueType& new_payload = values[i]; - // TODO(Alvaro) - Rethink this. I think it's fine for us to interpret empty values as a regular update - // (it'd empty out the payload of the zero leaf) if (new_payload.is_empty()) { continue; } + index_t index_of_new_leaf = i + meta.size; fr value = new_payload.get_key(); // This gives us the leaf that need updating @@ -1609,17 +1595,17 @@ void ContentAddressedIndexedTree::generate_sequential_inse .updated_leaf = IndexedLeafValueType::empty(), .original_leaf = low_leaf, }, - .new_leaf = std::nullopt, + .new_leaf = IndexedLeafValueType::empty(), + .new_leaf_index = index_of_new_leaf, }; if (!is_already_present) { // Update the current leaf to point it to the new leaf IndexedLeafValueType new_leaf = IndexedLeafValueType(new_payload, low_leaf.nextIndex, low_leaf.nextValue); - index_t index_of_new_leaf = current_size; + low_leaf.nextIndex = index_of_new_leaf; low_leaf.nextValue = value; - current_size++; // Cache the new leaf store_->set_leaf_key_at_index(index_of_new_leaf, new_leaf); store_->put_cached_leaf_by_index(index_of_new_leaf, new_leaf); @@ -1627,7 +1613,7 @@ void ContentAddressedIndexedTree::generate_sequential_inse store_->put_cached_leaf_by_index(low_leaf_index, low_leaf); insertion_update.low_leaf_update.updated_leaf = low_leaf; - insertion_update.new_leaf = std::pair(new_leaf, index_of_new_leaf); + insertion_update.new_leaf = new_leaf; } else if (IndexedLeafValueType::is_updateable()) { // Update the current leaf's value, don't change it's link IndexedLeafValueType replacement_leaf = @@ -1645,20 +1631,8 @@ void ContentAddressedIndexedTree::generate_sequential_inse " is already present")); } - response.inner.updates_to_perform.push_back(insertion_update); - } - - // Ensure that the tree is not going to be overfilled - if (current_size > max_size_) { - throw std::runtime_error(format("Unable to insert values into tree ", - meta.name, - " new size: ", - current_size, - " max size: ", - max_size_)); + response.inner.updates_to_perform->push_back(insertion_update); } - // The highest index touched will be current_size - 1 - response.inner.highest_index = current_size - 1; }, completion); } diff --git a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp index bc13a93e598..615f2ce4cf2 100644 --- a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp +++ b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/indexed_tree/indexed_leaf.hpp @@ -107,7 +107,7 @@ struct PublicDataLeafValue { fr get_key() const { return slot; } - bool is_empty() const { return slot == fr::zero() && value == fr::zero(); } + bool is_empty() const { return slot == fr::zero(); } std::vector get_hash_inputs(fr nextValue, fr nextIndex) const { diff --git a/barretenberg/cpp/src/barretenberg/vm/avm/trace/gadgets/merkle_tree.cpp b/barretenberg/cpp/src/barretenberg/vm/avm/trace/gadgets/merkle_tree.cpp index 32a1e7eec2b..00171489efa 100644 --- a/barretenberg/cpp/src/barretenberg/vm/avm/trace/gadgets/merkle_tree.cpp +++ b/barretenberg/cpp/src/barretenberg/vm/avm/trace/gadgets/merkle_tree.cpp @@ -102,8 +102,12 @@ FF AvmMerkleTreeTraceBuilder::perform_storage_write([[maybe_unused]] uint32_t cl // We update the low value low_preimage.value = value; FF low_preimage_hash = unconstrained_hash_public_data_preimage(low_preimage); - // Update the low leaf - return unconstrained_update_leaf_index(low_preimage_hash, static_cast(low_index), low_path); + // Update the low leaf - this will be returned in future + [[maybe_unused]] FF root = + unconstrained_update_leaf_index(low_preimage_hash, static_cast(low_index), low_path); + // TEMPORARY UNTIL WE CHANGE HOW UPDATES WORK + // Insert a zero leaf at the insertion index + return unconstrained_update_leaf_index(FF::zero(), static_cast(insertion_index), insertion_path); } // The new leaf for an insertion is PublicDataTreeLeafPreimage new_preimage{ diff --git a/noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/public_data_tree.nr b/noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/public_data_tree.nr index 4afaf9c75ce..e3d152ba98d 100644 --- a/noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/public_data_tree.nr +++ b/noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/components/public_data_tree.nr @@ -57,11 +57,16 @@ pub(crate) fn public_data_tree_insert( }, |write: PublicDataTreeLeaf, low_preimage: PublicDataTreeLeafPreimage| { // Build insertion leaf - PublicDataTreeLeafPreimage { - slot: write.slot, - value: write.value, - next_slot: low_preimage.next_slot, - next_index: low_preimage.next_index, + let is_update = low_preimage.slot == write.slot; + if is_update { + PublicDataTreeLeafPreimage::empty() + } else { + PublicDataTreeLeafPreimage { + slot: write.slot, + value: write.value, + next_slot: low_preimage.next_slot, + next_index: low_preimage.next_index, + } } }, ) diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/data/public_data_tree_leaf.nr b/noir-projects/noir-protocol-circuits/crates/types/src/data/public_data_tree_leaf.nr index 9cae7a6a675..e97a2161416 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/data/public_data_tree_leaf.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/data/public_data_tree_leaf.nr @@ -1,4 +1,4 @@ -use crate::{merkle_tree::leaf_preimage::IndexedTreeLeafValue, traits::Empty}; +use crate::traits::Empty; pub struct PublicDataTreeLeaf { pub slot: Field, @@ -17,12 +17,6 @@ impl Empty for PublicDataTreeLeaf { } } -impl IndexedTreeLeafValue for PublicDataTreeLeaf { - fn get_key(self) -> Field { - self.slot - } -} - impl PublicDataTreeLeaf { pub fn is_empty(self) -> bool { (self.slot == 0) & (self.value == 0) diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree.nr b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree.nr index 6b1be23c529..d6d2a363fd3 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree.nr @@ -3,7 +3,6 @@ pub mod check_valid_low_leaf; use crate::{ abis::append_only_tree_snapshot::AppendOnlyTreeSnapshot, merkle_tree::{ - leaf_preimage::{IndexedTreeLeafPreimage, IndexedTreeLeafValue}, membership::{assert_check_membership, MembershipWitness}, root::{calculate_empty_tree_root, calculate_subtree_root, root_from_sibling_path}, }, @@ -113,14 +112,14 @@ pub fn insert( build_insertion_leaf: fn(Value, Leaf) -> Leaf, ) -> AppendOnlyTreeSnapshot where - Value: IndexedTreeLeafValue, - Leaf: IndexedTreeLeafPreimage, + Value: Eq + Empty, + Leaf: Hash + Empty, { assert(is_valid_low_leaf(low_leaf_preimage, value), "Invalid low leaf"); // perform membership check for the low leaf against the original root assert_check_membership( - low_leaf_preimage.as_leaf(), + low_leaf_preimage.hash(), low_leaf_membership_witness.leaf_index, low_leaf_membership_witness.sibling_path, snapshot.root, @@ -130,34 +129,29 @@ where let updated_low_leaf = update_low_leaf(low_leaf_preimage, value, snapshot.next_available_leaf_index); - // Update low leaf snapshot.root = root_from_sibling_path( - updated_low_leaf.as_leaf(), + updated_low_leaf.hash(), low_leaf_membership_witness.leaf_index, low_leaf_membership_witness.sibling_path, ); - if low_leaf_preimage.get_key() == value.get_key() { - // If it's an update, we don't need to insert the new leaf and advance the tree - snapshot - } else { - let insertion_leaf = build_insertion_leaf(value, low_leaf_preimage); - assert_check_membership( - 0, - snapshot.next_available_leaf_index as Field, - insertion_sibling_path, - snapshot.root, - ); - - // Calculate the new root - snapshot.root = root_from_sibling_path( - insertion_leaf.as_leaf(), - snapshot.next_available_leaf_index as Field, - insertion_sibling_path, - ); - - snapshot.next_available_leaf_index += 1; - - snapshot - } + let insertion_leaf = build_insertion_leaf(value, low_leaf_preimage); + + assert_check_membership( + 0, + snapshot.next_available_leaf_index as Field, + insertion_sibling_path, + snapshot.root, + ); + + // Calculate the new root + snapshot.root = root_from_sibling_path( + insertion_leaf.hash(), + snapshot.next_available_leaf_index as Field, + insertion_sibling_path, + ); + + snapshot.next_available_leaf_index += 1; + + snapshot } diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree/check_valid_low_leaf.nr b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree/check_valid_low_leaf.nr index 9a42a21dafe..6c454c5f583 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree/check_valid_low_leaf.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/indexed_tree/check_valid_low_leaf.nr @@ -16,19 +16,12 @@ mod tests { indexed_tree::check_valid_low_leaf::assert_check_valid_low_leaf, leaf_preimage::IndexedTreeLeafPreimage, }; - use crate::traits::Empty; struct TestLeafPreimage { value: Field, next_value: Field, } - impl Empty for TestLeafPreimage { - fn empty() -> Self { - Self { value: 0, next_value: 0 } - } - } - impl IndexedTreeLeafPreimage for TestLeafPreimage { fn get_key(self) -> Field { self.value diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/leaf_preimage.nr b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/leaf_preimage.nr index dab825f5664..f33dd072d96 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/leaf_preimage.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/leaf_preimage.nr @@ -1,16 +1,10 @@ -use crate::traits::Empty; - pub trait LeafPreimage { fn get_key(self) -> Field; fn as_leaf(self) -> Field; } -pub trait IndexedTreeLeafPreimage: Empty { +pub trait IndexedTreeLeafPreimage { fn get_key(self) -> Field; fn get_next_key(self) -> Field; fn as_leaf(self) -> Field; } - -pub trait IndexedTreeLeafValue: Eq + Empty { - fn get_key(self) -> Field; -} diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/membership.nr b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/membership.nr index 34ab55a678d..178449af97d 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/membership.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/membership.nr @@ -93,7 +93,6 @@ mod tests { }, tests::merkle_tree_utils::NonEmptyMerkleTree, }; - use crate::traits::Empty; use std::hash::pedersen_hash; struct TestLeafPreimage { @@ -101,12 +100,6 @@ mod tests { next_value: Field, } - impl Empty for TestLeafPreimage { - fn empty() -> Self { - TestLeafPreimage { value: 0, next_value: 0 } - } - } - impl LeafPreimage for TestLeafPreimage { fn get_key(self) -> Field { self.value diff --git a/yarn-project/simulator/src/avm/avm_simulator.test.ts b/yarn-project/simulator/src/avm/avm_simulator.test.ts index 72889ea63c1..06811af1557 100644 --- a/yarn-project/simulator/src/avm/avm_simulator.test.ts +++ b/yarn-project/simulator/src/avm/avm_simulator.test.ts @@ -1178,7 +1178,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { const { preimage: lowLeafPreimage, index: lowLeafIndex, - alreadyPresent: leafAlreadyPresent, + update: leafAlreadyPresent, } = await ephemeralForest.getLeafOrLowLeafInfo( MerkleTreeId.PUBLIC_DATA_TREE, leafSlot0, @@ -1223,7 +1223,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { const { preimage: lowLeafPreimage, index: lowLeafIndex, - alreadyPresent: leafAlreadyPresent, + update: leafAlreadyPresent, } = await ephemeralForest.getLeafOrLowLeafInfo( MerkleTreeId.PUBLIC_DATA_TREE, leafSlot0, @@ -1294,7 +1294,7 @@ describe('AVM simulator: transpiled Noir contracts', () => { const { preimage: lowLeafPreimage, index: lowLeafIndex, - alreadyPresent: leafAlreadyPresent, + update: leafAlreadyPresent, } = await ephemeralForest.getLeafOrLowLeafInfo( MerkleTreeId.PUBLIC_DATA_TREE, leafSlot0, diff --git a/yarn-project/simulator/src/avm/avm_tree.test.ts b/yarn-project/simulator/src/avm/avm_tree.test.ts index 47f88232828..aa0ab6bc07b 100644 --- a/yarn-project/simulator/src/avm/avm_tree.test.ts +++ b/yarn-project/simulator/src/avm/avm_tree.test.ts @@ -1,15 +1,15 @@ import { + type BatchInsertionResult, type IndexedTreeId, MerkleTreeId, type MerkleTreeWriteOperations, - type SequentialInsertionResult, } from '@aztec/circuit-types'; import { NOTE_HASH_TREE_HEIGHT, NULLIFIER_SUBTREE_HEIGHT, - NULLIFIER_TREE_HEIGHT, + type NULLIFIER_TREE_HEIGHT, type NullifierLeafPreimage, - PUBLIC_DATA_TREE_HEIGHT, + type PUBLIC_DATA_TREE_HEIGHT, PublicDataTreeLeaf, type PublicDataTreeLeafPreimage, } from '@aztec/circuits.js'; @@ -18,16 +18,11 @@ import { Fr } from '@aztec/foundation/fields'; import { type IndexedTreeLeafPreimage } from '@aztec/foundation/trees'; import { openTmpStore } from '@aztec/kv-store/utils'; import { NoopTelemetryClient } from '@aztec/telemetry-client/noop'; -import { MerkleTrees, NativeWorldStateService } from '@aztec/world-state'; +import { MerkleTrees } from '@aztec/world-state'; -import { - AvmEphemeralForest, - EphemeralAvmTree, - type IndexedInsertResult, - type IndexedUpsertResult, -} from './avm_tree.js'; +import { AvmEphemeralForest, EphemeralAvmTree, type IndexedInsertionResult } from './avm_tree.js'; -let mainState: MerkleTreeWriteOperations; +let worldStateTrees: MerkleTrees; let copyState: MerkleTreeWriteOperations; // Up to 64 dummy note hashes let noteHashes: Fr[]; @@ -42,9 +37,8 @@ let getSiblingIndex = 21n; // Helper to check the equality of the insertion results (low witness, insertion path) const checkEqualityOfInsertionResults = ( - containerResults: IndexedUpsertResult[] | IndexedInsertResult[], - wsResults: SequentialInsertionResult[], - treeHeight: number, + containerResults: IndexedInsertionResult[], + wsResults: BatchInsertionResult[], ) => { if (containerResults.length !== wsResults.length) { throw new Error('Results length mismatch'); @@ -55,41 +49,40 @@ const checkEqualityOfInsertionResults = ( expect(containerResult.lowWitness.siblingPath).toEqual(wsResult.lowLeavesWitnessData![0].siblingPath.toFields()); expect(containerResult.lowWitness.index).toEqual(wsResult.lowLeavesWitnessData![0].index); expect(containerResult.lowWitness.preimage).toEqual(wsResult.lowLeavesWitnessData![0].leafPreimage); - if ('update' in containerResult && containerResult.update) { - expect(Array(treeHeight).fill(Fr.ZERO)).toEqual(wsResult.insertionWitnessData[0].siblingPath.toFields()); - } else { - expect(containerResult.insertionPath).toEqual(wsResult.insertionWitnessData[0].siblingPath.toFields()); - } + expect(containerResult.insertionPath).toEqual(wsResult.newSubtreeSiblingPath.toFields()); } }; const getWorldStateRoot = async (treeId: MerkleTreeId) => { - return (await mainState.getTreeInfo(treeId)).root; + return (await worldStateTrees.getTreeInfo(treeId, /*includeUncommitted=*/ true)).root; }; const getWorldStateSiblingPath = (treeId: MerkleTreeId, index: bigint) => { - return mainState.getSiblingPath(treeId, index); + return worldStateTrees.getSiblingPath(treeId, index, /*includeUncommitted=*/ true); }; const publicDataInsertWorldState = ( slot: Fr, value: Fr, -): Promise> => { - return mainState.sequentialInsert(MerkleTreeId.PUBLIC_DATA_TREE, [new PublicDataTreeLeaf(slot, value).toBuffer()]); +): Promise> => { + return worldStateTrees.batchInsert( + MerkleTreeId.PUBLIC_DATA_TREE, + [new PublicDataTreeLeaf(slot, value).toBuffer()], + 0, + ); }; const nullifierInsertWorldState = ( nullifier: Fr, -): Promise> => { - return mainState.sequentialInsert(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()]); +): Promise> => { + return worldStateTrees.batchInsert(MerkleTreeId.NULLIFIER_TREE, [nullifier.toBuffer()], 0); }; // Set up some recurring state for the tests beforeEach(async () => { - const worldState = await NativeWorldStateService.tmp(); - - mainState = await worldState.fork(); - copyState = await worldState.fork(); + const store = openTmpStore(true); + worldStateTrees = await MerkleTrees.new(store, new NoopTelemetryClient()); + copyState = await worldStateTrees.fork(); noteHashes = Array.from({ length: 64 }, (_, i) => new Fr(i)); // We do + 128 since the first 128 leaves are already filled in the indexed trees (nullifier, public data) @@ -132,7 +125,7 @@ describe('Simple Note Hash Consistency', () => { for (const noteHash of noteHashes) { treeContainer.appendNoteHash(noteHash); } - await mainState.appendLeaves(treeId, noteHashes); + await worldStateTrees.appendLeaves(treeId, noteHashes); // Check that the roots are consistent const wsRoot = await getWorldStateRoot(treeId); @@ -159,7 +152,7 @@ describe('Simple Note Hash Consistency', () => { } // Build a worldstateDB with all the note hashes - await mainState.appendLeaves(treeId, preInserted.concat(postInserted)); + await worldStateTrees.appendLeaves(treeId, preInserted.concat(postInserted)); // Check that the roots are consistent const wsRoot = await getWorldStateRoot(treeId); @@ -181,8 +174,8 @@ describe('Simple Note Hash Consistency', () => { describe('Simple Public Data Consistency', () => { const treeId = MerkleTreeId.PUBLIC_DATA_TREE as IndexedTreeId; - let containerInsertionResults: IndexedUpsertResult[] = []; - let worldStateInsertionResults: SequentialInsertionResult[] = []; + let containerInsertionResults: IndexedInsertionResult[] = []; + let worldStateInsertionResults: BatchInsertionResult[] = []; // We need to zero out between tests afterEach(() => { @@ -206,7 +199,7 @@ describe('Simple Public Data Consistency', () => { expect(computedRoot.toBuffer()).toEqual(wsRoot); // Check that all the accumulated insertion results match - checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults, PUBLIC_DATA_TREE_HEIGHT); + checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults); }); it('Should fork a prefilled tree and check consistency', async () => { @@ -274,14 +267,14 @@ describe('Simple Public Data Consistency', () => { expect(computedRoot.toBuffer()).toEqual(wsRoot); // Check the insertion results match - checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults, PUBLIC_DATA_TREE_HEIGHT); + checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults); }); }); describe('Simple Nullifier Consistency', () => { const treeId = MerkleTreeId.NULLIFIER_TREE as IndexedTreeId; - let containerInsertionResults: IndexedInsertResult[] = []; - let worldStateInsertionResults: SequentialInsertionResult[] = []; + let containerInsertionResults: IndexedInsertionResult[] = []; + let worldStateInsertionResults: BatchInsertionResult[] = []; // We need to zero out between tests afterEach(() => { @@ -304,7 +297,7 @@ describe('Simple Nullifier Consistency', () => { expect(computedRoot.toBuffer()).toEqual(wsRoot); // Check that all the accumulated insertion results match - checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults, NULLIFIER_TREE_HEIGHT); + checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults); // Check a sibling path from a random index is consistent const wsSiblingPath = await getWorldStateSiblingPath(treeId, getSiblingIndex); @@ -332,11 +325,7 @@ describe('Simple Nullifier Consistency', () => { expect(computedRoot.toBuffer()).toEqual(wsRoot); // Check insertion results - note we can only compare against the post-insertion results - checkEqualityOfInsertionResults( - containerInsertionResults, - worldStateInsertionResults.slice(preInsertIndex), - NULLIFIER_TREE_HEIGHT, - ); + checkEqualityOfInsertionResults(containerInsertionResults, worldStateInsertionResults.slice(preInsertIndex)); }); it('Should check that the insertion paths resolve to the root', async () => { @@ -377,7 +366,7 @@ describe('Simple Nullifier Consistency', () => { // Update low nullifier const newLowNullifier = containerInsert.lowWitness.preimage; newLowNullifier.nextIndex = containerInsert.leafIndex; - newLowNullifier.nextNullifier = containerInsert.element.nullifier; + newLowNullifier.nextNullifier = containerInsert.newOrElementToUpdate.element.nullifier; // Compute new root const updatedRoot = calcRootFromPath( containerInsert.lowWitness.siblingPath, @@ -392,7 +381,7 @@ describe('Simple Nullifier Consistency', () => { // Step 4 const finalRoot = calcRootFromPath( containerInsert.insertionPath, - treeContainer.hashPreimage(containerInsert.element), + treeContainer.hashPreimage(containerInsert.newOrElementToUpdate.element), containerInsert.leafIndex, ); expect(finalRoot.toBuffer()).toEqual(rootAfter); @@ -421,7 +410,7 @@ describe('Big Random Avm Ephemeral Container Test', () => { // Insert values ino merkleTrees // Note Hash - await mainState.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, noteHashes.slice(0, ENTRY_COUNT)); + await worldStateTrees.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, noteHashes.slice(0, ENTRY_COUNT)); // Everything else for (let i = 0; i < ENTRY_COUNT; i++) { await nullifierInsertWorldState(indexedHashes[i]); @@ -488,7 +477,7 @@ describe('Checking forking and merging', () => { it('Fork-Rollback-Fork-Merge should be consistent', async () => { // To store results - const wsInsertionResults: SequentialInsertionResult[] = []; + const wsInsertionResults: BatchInsertionResult[] = []; const containerInsertionResults = []; const treeContainer = await AvmEphemeralForest.create(copyState); @@ -517,7 +506,7 @@ describe('Checking forking and merging', () => { expect(containerRoot.toBuffer()).toEqual(wsRoot); // Check that all the accumulated insertion results - checkEqualityOfInsertionResults(containerInsertionResults, wsInsertionResults, PUBLIC_DATA_TREE_HEIGHT); + checkEqualityOfInsertionResults(containerInsertionResults, wsInsertionResults); }); }); @@ -565,7 +554,7 @@ describe('Batch Insertion', () => { const treeContainer = await AvmEphemeralForest.create(copyState); await treeContainer.appendNullifier(indexedHashes[0]); await treeContainer.appendNullifier(indexedHashes[1]); - await mainState.batchInsert( + await worldStateTrees.batchInsert( MerkleTreeId.NULLIFIER_TREE, [indexedHashes[0].toBuffer(), indexedHashes[1].toBuffer()], NULLIFIER_SUBTREE_HEIGHT, diff --git a/yarn-project/simulator/src/avm/avm_tree.ts b/yarn-project/simulator/src/avm/avm_tree.ts index f9cc70f745e..86c4b160d7d 100644 --- a/yarn-project/simulator/src/avm/avm_tree.ts +++ b/yarn-project/simulator/src/avm/avm_tree.ts @@ -17,14 +17,7 @@ import cloneDeep from 'lodash.clonedeep'; type PreimageWitness = { preimage: T; index: bigint; -}; - -/** - * The result of fetching a leaf from an indexed tree. Contains the preimage and wether the leaf was already present - * or it's a low leaf. - */ -type GetLeafResult = PreimageWitness & { - alreadyPresent: boolean; + update: boolean; }; /** @@ -36,30 +29,16 @@ type LeafWitness = PreimageWitness & { }; /** - * The result of an update in an indexed merkle tree (no new leaf inserted) - */ -type IndexedUpdateResult = { - element: T; - lowWitness: LeafWitness; -}; - -/** - * The result of an insertion in an indexed merkle tree. + * The result of an indexed insertion in an indexed merkle tree. * This will be used to hint the circuit */ -export type IndexedInsertResult = IndexedUpdateResult & { +export type IndexedInsertionResult = { leafIndex: bigint; insertionPath: Fr[]; + newOrElementToUpdate: { update: boolean; element: T }; + lowWitness: LeafWitness; }; -/** - * The result of an indexed upsert in an indexed merkle tree. - * This will be used to hint the circuit - */ -export type IndexedUpsertResult = - | (IndexedUpdateResult & { update: true }) - | (IndexedInsertResult & { update: false }); - /****************************************************/ /****** The AvmEphemeralForest Class ****************/ /****************************************************/ @@ -165,7 +144,7 @@ export class AvmEphemeralForest { * @param newValue - The value to be written or updated to * @returns The insertion result which contains the insertion path, low leaf and the new leaf index */ - async writePublicStorage(slot: Fr, newValue: Fr): Promise> { + async writePublicStorage(slot: Fr, newValue: Fr): Promise> { // This only works for the public data tree const treeId = MerkleTreeId.PUBLIC_DATA_TREE; const tree = this.treeMap.get(treeId)!; @@ -173,12 +152,12 @@ export class AvmEphemeralForest { typeof treeId, PublicDataTreeLeafPreimage >(treeId, slot); - const { preimage, index: lowLeafIndex, alreadyPresent: update } = leafOrLowLeafInfo; - const siblingPath = await this.getSiblingPath(treeId, lowLeafIndex); + const { preimage, index, update } = leafOrLowLeafInfo; + const siblingPath = await this.getSiblingPath(treeId, index); if (pathAbsentInEphemeralTree) { // Since we have never seen this before - we should insert it into our tree as it is about to be updated. - this.treeMap.get(treeId)!.insertSiblingPath(lowLeafIndex, siblingPath); + this.treeMap.get(treeId)!.insertSiblingPath(index, siblingPath); } if (update) { @@ -186,18 +165,29 @@ export class AvmEphemeralForest { const existingPublicDataSiblingPath = siblingPath; updatedPreimage.value = newValue; - tree.updateLeaf(this.hashPreimage(updatedPreimage), lowLeafIndex); - this.setIndexedUpdates(treeId, lowLeafIndex, updatedPreimage); - this._updateSortedKeys(treeId, [updatedPreimage.slot], [lowLeafIndex]); + // It is really unintuitive that by updating, we are also appending a Zero Leaf to the tree + // Additionally, this leaf preimage does not seem to factor into further appends + const emptyLeaf = new PublicDataTreeLeafPreimage(Fr.ZERO, Fr.ZERO, Fr.ZERO, 0n); + const insertionIndex = tree.leafCount; + tree.updateLeaf(this.hashPreimage(updatedPreimage), index); + tree.appendLeaf(Fr.ZERO); + this.setIndexedUpdates(treeId, index, updatedPreimage); + this.setIndexedUpdates(treeId, insertionIndex, emptyLeaf); + const insertionPath = tree.getSiblingPath(insertionIndex)!; + + // Even though we append an empty leaf into the tree as a part of update - it doesnt seem to impact future inserts... + this._updateSortedKeys(treeId, [updatedPreimage.slot], [index]); return { - element: updatedPreimage, + leafIndex: insertionIndex, + insertionPath, + newOrElementToUpdate: { update: true, element: updatedPreimage }, lowWitness: { preimage: preimage, - index: lowLeafIndex, + index: index, + update: true, siblingPath: existingPublicDataSiblingPath, }, - update: true, }; } // We are writing to a new slot, so our preimage is a lowNullifier @@ -212,22 +202,22 @@ export class AvmEphemeralForest { new Fr(preimage.getNextKey()), preimage.getNextIndex(), ); - const insertionPath = this.appendIndexedTree(treeId, lowLeafIndex, updatedLowLeaf, newPublicDataLeaf); + const insertionPath = this.appendIndexedTree(treeId, index, updatedLowLeaf, newPublicDataLeaf); // Even though the low leaf key is not updated, we still need to update the sorted keys in case we have // not seen the low leaf before - this._updateSortedKeys(treeId, [newPublicDataLeaf.slot, updatedLowLeaf.slot], [insertionIndex, lowLeafIndex]); + this._updateSortedKeys(treeId, [newPublicDataLeaf.slot, updatedLowLeaf.slot], [insertionIndex, index]); return { - element: newPublicDataLeaf, + leafIndex: insertionIndex, + insertionPath: insertionPath, + newOrElementToUpdate: { update: false, element: newPublicDataLeaf }, lowWitness: { preimage, - index: lowLeafIndex, + index: index, + update: false, siblingPath, }, - update: false, - leafIndex: insertionIndex, - insertionPath: insertionPath, }; } @@ -257,14 +247,14 @@ export class AvmEphemeralForest { * @param value - The nullifier to be appended * @returns The insertion result which contains the insertion path, low leaf and the new leaf index */ - async appendNullifier(nullifier: Fr): Promise> { + async appendNullifier(nullifier: Fr): Promise> { const treeId = MerkleTreeId.NULLIFIER_TREE; const tree = this.treeMap.get(treeId)!; const [leafOrLowLeafInfo, pathAbsentInEphemeralTree] = await this._getLeafOrLowLeafInfo< typeof treeId, NullifierLeafPreimage >(treeId, nullifier); - const { preimage, index, alreadyPresent } = leafOrLowLeafInfo; + const { preimage, index, update } = leafOrLowLeafInfo; const siblingPath = await this.getSiblingPath(treeId, index); if (pathAbsentInEphemeralTree) { @@ -272,7 +262,7 @@ export class AvmEphemeralForest { this.treeMap.get(treeId)!.insertSiblingPath(index, siblingPath); } - assert(!alreadyPresent, 'Nullifier already exists in the tree. Cannot update a nullifier!'); + assert(!update, 'Nullifier already exists in the tree. Cannot update a nullifier!'); // We are writing a new entry const insertionIndex = tree.leafCount; @@ -292,14 +282,15 @@ export class AvmEphemeralForest { ); return { - element: newNullifierLeaf, + leafIndex: insertionIndex, + insertionPath: insertionPath, + newOrElementToUpdate: { update: false, element: newNullifierLeaf }, lowWitness: { preimage, index, + update, siblingPath, }, - leafIndex: insertionIndex, - insertionPath: insertionPath, }; } @@ -369,7 +360,7 @@ export class AvmEphemeralForest { async getLeafOrLowLeafInfo( treeId: ID, key: Fr, - ): Promise> { + ): Promise> { const [leafOrLowLeafInfo, _] = await this._getLeafOrLowLeafInfo(treeId, key); return leafOrLowLeafInfo; } @@ -382,14 +373,14 @@ export class AvmEphemeralForest { * @param key - The key for which we are look up the leaf or low leaf for. * @param T - The type of the preimage (PublicData or Nullifier) * @returns [ - * getLeafResult - The leaf or low leaf info (preimage & leaf index), + * preimageWitness - The leaf or low leaf info (preimage & leaf index), * pathAbsentInEphemeralTree - whether its sibling path is absent in the ephemeral tree (useful during insertions) * ] */ async _getLeafOrLowLeafInfo( treeId: ID, key: Fr, - ): Promise<[GetLeafResult, /*pathAbsentInEphemeralTree=*/ boolean]> { + ): Promise<[PreimageWitness, /*pathAbsentInEphemeralTree=*/ boolean]> { const bigIntKey = key.toBigInt(); // In this function, "min" refers to the leaf with the // largest key <= the specified key in the indexedUpdates. @@ -401,7 +392,7 @@ export class AvmEphemeralForest { if (minIndexedLeafIndex === -1n) { // No leaf is present in the indexed updates that is <= the key, // so retrieve the leaf or low leaf from the underlying DB. - const leafOrLowLeafPreimage: GetLeafResult = await this._getLeafOrLowLeafWitnessInExternalDb( + const leafOrLowLeafPreimage: PreimageWitness = await this._getLeafOrLowLeafWitnessInExternalDb( treeId, bigIntKey, ); @@ -411,7 +402,7 @@ export class AvmEphemeralForest { const minPreimage: T = this.getIndexedUpdate(treeId, minIndexedLeafIndex); if (minPreimage.getKey() === bigIntKey) { // the index found is an exact match, no need to search further - const leafInfo = { preimage: minPreimage, index: minIndexedLeafIndex, alreadyPresent: true }; + const leafInfo = { preimage: minPreimage, index: minIndexedLeafIndex, update: true }; return [leafInfo, /*pathAbsentInEphemeralTree=*/ false]; } else { // We are starting with the leaf with largest key <= the specified key @@ -462,7 +453,7 @@ export class AvmEphemeralForest { private async _getLeafOrLowLeafWitnessInExternalDb( treeId: ID, key: bigint, - ): Promise> { + ): Promise> { // "key" is siloed slot (leafSlot) or siloed nullifier const previousValueIndex = await this.treeDb.getPreviousValueIndex(treeId, key); assert( @@ -477,7 +468,7 @@ export class AvmEphemeralForest { `${MerkleTreeId[treeId]} low leaf preimage should never be undefined (even if target leaf does not exist)`, ); - return { preimage: leafPreimage as T, index: leafIndex, alreadyPresent }; + return { preimage: leafPreimage as T, index: leafIndex, update: alreadyPresent }; } /** @@ -490,7 +481,7 @@ export class AvmEphemeralForest { * @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 [ - * GetLeafResult | undefined - The leaf or low leaf info (preimage & leaf index), + * 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) * ] * @@ -506,10 +497,10 @@ export class AvmEphemeralForest { key: bigint, minPreimage: T, minIndex: bigint, - ): Promise<[GetLeafResult | undefined, /*pathAbsentInEphemeralTree=*/ boolean]> { + ): Promise<[PreimageWitness | undefined, /*pathAbsentInEphemeralTree=*/ boolean]> { let found = false; let curr = minPreimage as T; - let result: GetLeafResult | undefined = undefined; + let result: PreimageWitness | 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; @@ -520,11 +511,11 @@ export class AvmEphemeralForest { if (curr.getKey() === bigIntKey) { // We found an exact match - therefore this is an update found = true; - result = { preimage: curr, index: lowPublicDataIndex, alreadyPresent: true }; + result = { preimage: curr, index: lowPublicDataIndex, update: true }; } else if (curr.getKey() < bigIntKey && (curr.getNextIndex() === 0n || curr.getNextKey() > bigIntKey)) { // We found it via sandwich or max condition, this is a low nullifier found = true; - result = { preimage: curr, index: lowPublicDataIndex, alreadyPresent: false }; + result = { preimage: curr, index: lowPublicDataIndex, update: false }; } // Update the the values for the next iteration else { diff --git a/yarn-project/simulator/src/avm/journal/journal.ts b/yarn-project/simulator/src/avm/journal/journal.ts index 63dbf59f09e..9a3ffa5273a 100644 --- a/yarn-project/simulator/src/avm/journal/journal.ts +++ b/yarn-project/simulator/src/avm/journal/journal.ts @@ -163,12 +163,8 @@ export class AvmPersistableStateManager { const lowLeafIndex = lowLeafInfo.index; const lowLeafPath = lowLeafInfo.siblingPath; - const newLeafPreimage = result.element as PublicDataTreeLeafPreimage; - let insertionPath; - - if (!result.update) { - insertionPath = result.insertionPath; - } + const insertionPath = result.insertionPath; + const newLeafPreimage = result.newOrElementToUpdate.element as PublicDataTreeLeafPreimage; this.trace.tracePublicStorageWrite( contractAddress, @@ -204,7 +200,7 @@ export class AvmPersistableStateManager { const { preimage, index: leafIndex, - alreadyPresent, + update: exists, } = await this.merkleTrees.getLeafOrLowLeafInfo(MerkleTreeId.PUBLIC_DATA_TREE, leafSlot); // The index and preimage here is either the low leaf or the leaf itself (depending on the value of update flag) // In either case, we just want the sibling path to this leaf - it's up to the avm to distinguish if it's a low leaf or not @@ -216,7 +212,7 @@ export class AvmPersistableStateManager { ); this.log.debug(`leafPreimage.slot: ${leafPreimage.slot}, leafPreimage.value: ${leafPreimage.value}`); - if (!alreadyPresent) { + if (!exists) { // Sanity check that the leaf slot is skipped by low leaf when it doesn't exist assert( leafPreimage.slot.toBigInt() < leafSlot.toBigInt() && @@ -312,15 +308,12 @@ export class AvmPersistableStateManager { const { preimage, index: leafIndex, - alreadyPresent, + update, } = await this.merkleTrees.getLeafOrLowLeafInfo(MerkleTreeId.NULLIFIER_TREE, siloedNullifier); const leafPreimage = preimage as NullifierLeafPreimage; const leafPath = await this.merkleTrees.getSiblingPath(MerkleTreeId.NULLIFIER_TREE, leafIndex); - assert( - alreadyPresent == exists, - 'WorldStateDB contains nullifier leaf, but merkle tree does not.... This is a bug!', - ); + assert(update == exists, 'WorldStateDB contains nullifier leaf, but merkle tree does not.... This is a bug!'); if (exists) { this.log.debug(`Siloed nullifier ${siloedNullifier} exists at leafIndex=${leafIndex}`); @@ -362,11 +355,11 @@ export class AvmPersistableStateManager { // Maybe overkill, but we should check if the nullifier is already present in the tree before attempting to insert // It might be better to catch the error from the insert operation // Trace all nullifier creations, even duplicate insertions that fail - const { preimage, index, alreadyPresent } = await this.merkleTrees.getLeafOrLowLeafInfo( + const { preimage, index, update } = await this.merkleTrees.getLeafOrLowLeafInfo( MerkleTreeId.NULLIFIER_TREE, siloedNullifier, ); - if (alreadyPresent) { + if (update) { this.log.verbose(`Siloed nullifier ${siloedNullifier} already present in tree at index ${index}!`); // If the nullifier is already present, we should not insert it again // instead we provide the direct membership path @@ -374,7 +367,7 @@ export class AvmPersistableStateManager { // This just becomes a nullifier read hint this.trace.traceNullifierCheck( siloedNullifier, - /*exists=*/ alreadyPresent, + /*exists=*/ update, preimage as NullifierLeafPreimage, new Fr(index), path,