feat: Add tree equality assertions (#10756)

Asserts in the public base that the end tree roots of the AVM equal the
calculated tree roots in the public base, to check correctness of the
AVM insertions.

barretenberg/cpp/src/barretenberg/vm/avm/tests/execution.test.cpp

@@ -54,6 +54,10 @@ class AvmExecutionTests : public ::testing::Test {
         public_inputs.gas_settings.gas_limits.da_gas = DEFAULT_INITIAL_DA_GAS;
         public_inputs.start_gas_used.l2_gas = 0;
         public_inputs.start_gas_used.da_gas = 0;
+        public_inputs.end_tree_snapshots.note_hash_tree.size =
+            public_inputs.start_tree_snapshots.note_hash_tree.size + MAX_NOTE_HASHES_PER_TX;
+        public_inputs.end_tree_snapshots.nullifier_tree.size =
+            public_inputs.start_tree_snapshots.nullifier_tree.size + MAX_NULLIFIERS_PER_TX;
 
         // These values are magic because of how some tests work! Don't change them
         PublicCallRequest dummy_request = {

barretenberg/cpp/src/barretenberg/vm/avm/trace/execution.cpp

@@ -399,6 +399,8 @@ std::vector<Row> Execution::gen_trace(AvmPublicInputs const& public_inputs,
         trace_builder.pay_fee();
     }
 
+    trace_builder.pad_trees();
+
     auto trace = trace_builder.finalize(apply_e2e_assertions);
 
     returndata = trace_builder.get_all_returndata();

barretenberg/cpp/src/barretenberg/vm/avm/trace/gadgets/merkle_tree.hpp

@@ -69,6 +69,7 @@ class AvmMerkleTreeTraceBuilder {
                                 const std::vector<FF>& low_path,
                                 const FF& nullifier,
                                 const std::vector<FF>& insertion_path);
+    void set_nullifier_tree_size(uint32_t size) { tree_snapshots.nullifier_tree.size = size; }
 
     // Note Hash Tree
     bool perform_note_hash_read(uint32_t clk,
@@ -77,6 +78,7 @@ class AvmMerkleTreeTraceBuilder {
                                 const std::vector<FF>& path) const;
 
     FF perform_note_hash_append(uint32_t clk, const FF& note_hash, const std::vector<FF>& insertion_path);
+    void set_note_hash_tree_size(uint32_t size) { tree_snapshots.note_hash_tree.size = size; }
 
     // L1 to L2 Message Tree
     bool perform_l1_to_l2_message_read(uint32_t clk,

barretenberg/cpp/src/barretenberg/vm/avm/trace/trace.cpp

@@ -312,6 +312,23 @@ void AvmTraceBuilder::pay_fee()
     side_effect_counter++;
 }
 
+void AvmTraceBuilder::pad_trees()
+{
+    auto current_tree_snapshots = merkle_tree_trace_builder.get_tree_snapshots();
+
+    auto initial_note_hash_snapshot = public_inputs.start_tree_snapshots.note_hash_tree;
+    // sanity check
+    uint32_t padded_note_hash_size = initial_note_hash_snapshot.size + MAX_NOTE_HASHES_PER_TX;
+    ASSERT(current_tree_snapshots.note_hash_tree.size <= padded_note_hash_size);
+    merkle_tree_trace_builder.set_note_hash_tree_size(padded_note_hash_size);
+
+    auto initial_nullifier_snapshot = public_inputs.start_tree_snapshots.nullifier_tree;
+    // sanity check
+    uint32_t padded_nullifier_size = initial_nullifier_snapshot.size + MAX_NULLIFIERS_PER_TX;
+    ASSERT(current_tree_snapshots.nullifier_tree.size <= padded_nullifier_size);
+    merkle_tree_trace_builder.set_nullifier_tree_size(padded_nullifier_size);
+}
+
 /**
  * @brief Loads a value from memory into a given intermediate register at a specified clock cycle.
  * Handles both direct and indirect memory access.

barretenberg/cpp/src/barretenberg/vm/avm/trace/trace.hpp

@@ -238,6 +238,7 @@ class AvmTraceBuilder {
     void insert_private_revertible_state(const std::vector<FF>& siloed_nullifiers,
                                          const std::vector<FF>& siloed_note_hashes);
     void pay_fee();
+    void pad_trees();
 
     // These are used for testing only.
     AvmTraceBuilder& set_range_check_required(bool required)

noir-projects/noir-protocol-circuits/crates/rollup-lib/src/base/public_base_rollup.nr

@@ -150,7 +150,6 @@ impl PublicBaseRollupInputs {
         // Validate public data update requests and update public data tree
         let end_public_data_tree_snapshot =
             self.validate_and_process_public_state(combined_accumulated_data.public_data_writes);
-
         // Append the tx effects for blob(s)
         let siloed_l2_to_l1_msgs = combined_accumulated_data.l2_to_l1_msgs.map(
             |message: ScopedL2ToL1Message| silo_l2_to_l1_message(
@@ -179,6 +178,25 @@ impl PublicBaseRollupInputs {
             combined_constant_data.historical_header,
         );
 
+        assert(
+            end_note_hash_tree_snapshot.eq(
+                self.avm_proof_data.public_inputs.end_tree_snapshots.note_hash_tree,
+            ),
+            "Mismatch note hash tree base rollup vs AVM",
+        );
+        assert(
+            end_nullifier_tree_snapshot.eq(
+                self.avm_proof_data.public_inputs.end_tree_snapshots.nullifier_tree,
+            ),
+            "Mismatch nullifier tree base rollup vs AVM",
+        );
+        assert(
+            end_public_data_tree_snapshot.eq(
+                self.avm_proof_data.public_inputs.end_tree_snapshots.public_data_tree,
+            ),
+            "Mismatch public data tree base rollup vs AVM",
+        );
+
         BaseOrMergeRollupPublicInputs {
             rollup_type: BASE_ROLLUP_TYPE,
             num_txs: 1,
@@ -297,6 +315,7 @@ mod tests {
             fixture_builder::FixtureBuilder,
             fixtures::{self, merkle_tree::generate_full_sha_tree},
             merkle_tree_utils::NonEmptyMerkleTree,
+            utils::pad_end,
         },
         traits::{Empty, is_empty},
         utils::{
@@ -320,7 +339,7 @@ mod tests {
         snapshot: AppendOnlyTreeSnapshot,
         writes: [(u32, PublicDataTreeLeaf); MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX],
         mut pre_existing_public_data: [PublicDataTreeLeafPreimage; EXISTING_LEAVES],
-    ) -> ([PublicDataTreeLeafPreimage; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX], [MembershipWitness<PUBLIC_DATA_TREE_HEIGHT>; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX], [[Field; PUBLIC_DATA_TREE_HEIGHT]; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX]) {
+    ) -> ([PublicDataTreeLeafPreimage; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX], [MembershipWitness<PUBLIC_DATA_TREE_HEIGHT>; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX], [[Field; PUBLIC_DATA_TREE_HEIGHT]; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX], AppendOnlyTreeSnapshot) {
         let mut low_leaves =
             [PublicDataTreeLeafPreimage::empty(); MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX];
         let mut low_public_data_writes_witnesses =
@@ -329,7 +348,6 @@ mod tests {
             [[0; PUBLIC_DATA_TREE_HEIGHT]; MAX_TOTAL_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX];
 
         let mut current_next_leaf_index = snapshot.next_available_leaf_index;
-
         for i in 0..writes.len() {
             let (low_leaf_index, write) = writes[i];
             if (!is_empty(write)) {
@@ -368,10 +386,18 @@ mod tests {
                 low_leaves[i] = low_leaf;
                 low_public_data_writes_witnesses[i] = low_public_data_writes_witness;
                 insertion_witnesses[i] = insertion_witness;
-                current_next_leaf_index += 1;
+                if low_leaf.slot != write.slot {
+                    current_next_leaf_index += 1;
+                }
             }
         }
-        (low_leaves, low_public_data_writes_witnesses, insertion_witnesses)
+        (
+            low_leaves, low_public_data_writes_witnesses, insertion_witnesses,
+            AppendOnlyTreeSnapshot {
+                root: public_data_tree.get_root(),
+                next_available_leaf_index: current_next_leaf_index,
+            },
+        )
     }
 
     struct PublicBaseRollupInputsBuilder {
@@ -445,11 +471,11 @@ mod tests {
             sibling_path
         }
 
-        fn update_nullifier_tree_with_new_leaves(
+        fn update_nullifier_tree_with_new_leaves<let AVAILABLE_LEAVES: u32, let TEST_SUPERTREE_HEIGHT: u32, let TEST_SUBTREE_HEIGHT: u32>(
             mut self,
-            nullifier_tree: &mut NonEmptyMerkleTree<MAX_NULLIFIERS_PER_TX, NULLIFIER_TREE_HEIGHT, NULLIFIER_SUBTREE_SIBLING_PATH_LENGTH, NULLIFIER_SUBTREE_HEIGHT>,
+            nullifier_tree: &mut NonEmptyMerkleTree<AVAILABLE_LEAVES, NULLIFIER_TREE_HEIGHT, TEST_SUPERTREE_HEIGHT, TEST_SUBTREE_HEIGHT>,
             start_nullifier_tree_snapshot: AppendOnlyTreeSnapshot,
-        ) -> ([NullifierLeafPreimage; MAX_NULLIFIERS_PER_TX], [MembershipWitness<NULLIFIER_TREE_HEIGHT>; MAX_NULLIFIERS_PER_TX], [Field; MAX_NULLIFIERS_PER_TX], [u32; MAX_NULLIFIERS_PER_TX]) {
+        ) -> ([NullifierLeafPreimage; MAX_NULLIFIERS_PER_TX], [MembershipWitness<NULLIFIER_TREE_HEIGHT>; MAX_NULLIFIERS_PER_TX], [Field; MAX_NULLIFIERS_PER_TX], [u32; MAX_NULLIFIERS_PER_TX], [Field; NULLIFIER_SUBTREE_SIBLING_PATH_LENGTH], AppendOnlyTreeSnapshot) {
             let mut nullifier_predecessor_preimages =
                 [NullifierLeafPreimage::empty(); MAX_NULLIFIERS_PER_TX];
             let mut low_nullifier_membership_witness =
@@ -476,6 +502,7 @@ mod tests {
             }
 
             let mut pre_existing_nullifiers = self.pre_existing_nullifiers;
+            let mut insertion_subtree = [NullifierLeafPreimage::empty(); MAX_NULLIFIERS_PER_TX];
 
             for i in 0..MAX_NULLIFIERS_PER_TEST {
                 if i < self.nullifiers.len() {
@@ -486,25 +513,50 @@ mod tests {
                     let low_index = self.nullifiers.get_unchecked(original_index).existing_index;
 
                     let mut low_preimage = pre_existing_nullifiers[low_index];
+                    let new_leaf = NullifierLeafPreimage {
+                        nullifier: new_nullifier,
+                        next_nullifier: low_preimage.next_nullifier,
+                        next_index: low_preimage.next_index,
+                    };
                     nullifier_predecessor_preimages[i] = low_preimage;
                     low_nullifier_membership_witness[i] = MembershipWitness {
                         leaf_index: low_index as Field,
                         sibling_path: nullifier_tree.get_sibling_path(low_index),
                     };
 
                     low_preimage.next_nullifier = new_nullifier;
-                    low_preimage.next_index = start_nullifier_tree_snapshot
-                        .next_available_leaf_index as u32
-                        + original_index;
+                    low_preimage.next_index =
+                        start_nullifier_tree_snapshot.next_available_leaf_index + original_index;
                     pre_existing_nullifiers[low_index] = low_preimage;
 
                     nullifier_tree.update_leaf(low_index, low_preimage.hash());
+                    insertion_subtree[original_index] = new_leaf;
                 }
             }
 
+            let nullifier_subtree_sibling_path = PublicBaseRollupInputsBuilder::extract_subtree_sibling_path(
+                nullifier_tree.get_sibling_path(self.pre_existing_nullifiers.len()),
+                [0; NULLIFIER_SUBTREE_SIBLING_PATH_LENGTH],
+            );
+
+            for i in 0..insertion_subtree.len() {
+                let leaf = insertion_subtree[i].as_leaf();
+                nullifier_tree.update_leaf(
+                    start_nullifier_tree_snapshot.next_available_leaf_index + i,
+                    leaf,
+                );
+            }
+
+            let end_nullifier_tree_snapshot = AppendOnlyTreeSnapshot {
+                root: nullifier_tree.get_root(),
+                next_available_leaf_index: start_nullifier_tree_snapshot.next_available_leaf_index
+                    + MAX_NULLIFIERS_PER_TX,
+            };
+
             (
                 nullifier_predecessor_preimages, low_nullifier_membership_witness,
-                sorted_nullifiers, sorted_nullifiers_indexes,
+                sorted_nullifiers, sorted_nullifiers_indexes, nullifier_subtree_sibling_path,
+                end_nullifier_tree_snapshot,
             )
         }
 
@@ -514,31 +566,47 @@ mod tests {
 
             avm_proof_data.public_inputs.transaction_fee = self.transaction_fee;
 
-            let start_note_hash_tree = NonEmptyMerkleTree::new(
-                self.pre_existing_notes,
+            let mut start_note_hash_tree = NonEmptyMerkleTree::new(
+                pad_end::<Field, MAX_NOTE_HASHES_PER_TX, MAX_NOTE_HASHES_PER_TX * 2>(
+                    self.pre_existing_notes,
+                    0,
+                ),
                 [0; NOTE_HASH_TREE_HEIGHT],
-                [0; NOTE_HASH_TREE_HEIGHT - NOTE_HASH_SUBTREE_HEIGHT],
-                [0; NOTE_HASH_SUBTREE_HEIGHT],
+                [0; NOTE_HASH_TREE_HEIGHT - NOTE_HASH_SUBTREE_HEIGHT - 1],
+                [0; NOTE_HASH_SUBTREE_HEIGHT + 1],
             );
             let start_note_hash_tree_snapshot = AppendOnlyTreeSnapshot {
                 root: start_note_hash_tree.get_root(),
-                next_available_leaf_index: start_note_hash_tree.get_next_available_index() as u32,
+                next_available_leaf_index: self.pre_existing_notes.len(),
             };
             let note_hash_subtree_sibling_path = PublicBaseRollupInputsBuilder::extract_subtree_sibling_path(
                 start_note_hash_tree.get_sibling_path(self.pre_existing_notes.len()),
                 [0; NOTE_HASH_SUBTREE_SIBLING_PATH_LENGTH],
             );
+            for i in 0..MAX_NOTE_HASHES_PER_TX {
+                let note_hash = avm_proof_data.public_inputs.accumulated_data.note_hashes[i];
+                start_note_hash_tree.update_leaf(MAX_NOTE_HASHES_PER_TX + i, note_hash);
+            }
+            avm_proof_data.public_inputs.end_tree_snapshots.note_hash_tree = AppendOnlyTreeSnapshot {
+                root: start_note_hash_tree.get_root(),
+                next_available_leaf_index: MAX_NOTE_HASHES_PER_TX * 2,
+            };
 
             let mut start_nullifier_tree = NonEmptyMerkleTree::new(
-                self.pre_existing_nullifiers.map(|preimage: NullifierLeafPreimage| preimage.hash()),
+                pad_end::<Field, MAX_NULLIFIERS_PER_TX, MAX_NULLIFIERS_PER_TX * 2>(
+                    self.pre_existing_nullifiers.map(|preimage: NullifierLeafPreimage| {
+                        preimage.hash()
+                    }),
+                    0,
+                ),
                 [0; NULLIFIER_TREE_HEIGHT],
-                [0; NULLIFIER_TREE_HEIGHT - NULLIFIER_SUBTREE_HEIGHT],
-                [0; NULLIFIER_SUBTREE_HEIGHT],
+                [0; NULLIFIER_TREE_HEIGHT - NULLIFIER_SUBTREE_HEIGHT - 1],
+                [0; NULLIFIER_SUBTREE_HEIGHT + 1],
             );
 
             let start_nullifier_tree_snapshot = AppendOnlyTreeSnapshot {
                 root: start_nullifier_tree.get_root(),
-                next_available_leaf_index: start_nullifier_tree.get_next_available_index() as u32,
+                next_available_leaf_index: MAX_NULLIFIERS_PER_TX,
             };
 
             let mut pre_existing_leaves = [0; AVAILABLE_PUBLIC_DATA_LEAVES_FOR_TEST];
@@ -569,7 +637,7 @@ mod tests {
                 next_available_leaf_index: start_archive.get_next_available_index() as u32,
             };
 
-            let (nullifier_predecessor_preimages, nullifier_predecessor_membership_witnesses, sorted_nullifiers, sorted_nullifier_indexes) = self
+            let (nullifier_predecessor_preimages, nullifier_predecessor_membership_witnesses, sorted_nullifiers, sorted_nullifier_indexes, nullifier_subtree_sibling_path, end_nullifier_tree_snapshot) = self
                 .update_nullifier_tree_with_new_leaves(
                     &mut start_nullifier_tree,
                     start_nullifier_tree_snapshot,
@@ -580,18 +648,19 @@ mod tests {
                 avm_proof_data.public_inputs.accumulated_data.nullifiers[i] = nullifier.value;
             }
 
-            let nullifier_subtree_sibling_path = PublicBaseRollupInputsBuilder::extract_subtree_sibling_path(
-                start_nullifier_tree.get_sibling_path(self.pre_existing_nullifiers.len()),
-                [0; NULLIFIER_SUBTREE_SIBLING_PATH_LENGTH],
-            );
+            avm_proof_data.public_inputs.end_tree_snapshots.nullifier_tree =
+                end_nullifier_tree_snapshot;
 
-            let (low_public_data_writes_preimages, low_public_data_writes_witnesses, public_data_tree_sibling_paths) = update_public_data_tree(
+            let (low_public_data_writes_preimages, low_public_data_writes_witnesses, public_data_tree_sibling_paths, end_public_data_tree_snapshot) = update_public_data_tree(
                 &mut start_public_data_tree,
                 start_public_data_tree_snapshot,
                 self.public_data_writes.storage(),
                 self.pre_existing_public_data,
             );
 
+            avm_proof_data.public_inputs.end_tree_snapshots.public_data_tree =
+                end_public_data_tree_snapshot;
+
             for i in 0..self.public_data_writes.len() {
                 let leaf = self.public_data_writes.get_unchecked(i).1;
                 avm_proof_data.public_inputs.accumulated_data.public_data_writes[i] =

yarn-project/prover-client/src/mocks/test_context.ts

@@ -8,7 +8,13 @@ import {
   type TxValidator,
 } from '@aztec/circuit-types';
 import { makeBloatedProcessedTx } from '@aztec/circuit-types/test';
-import { type AppendOnlyTreeSnapshot, BlockHeader, type Gas, type GlobalVariables } from '@aztec/circuits.js';
+import {
+  type AppendOnlyTreeSnapshot,
+  BlockHeader,
+  type Gas,
+  type GlobalVariables,
+  TreeSnapshots,
+} from '@aztec/circuits.js';
 import { times } from '@aztec/foundation/collection';
 import { Fr } from '@aztec/foundation/fields';
 import { type Logger } from '@aztec/foundation/log';
@@ -36,7 +42,7 @@ import { buildBlock } from '../block_builder/light.js';
 import { ProvingOrchestrator } from '../orchestrator/index.js';
 import { MemoryProvingQueue } from '../prover-agent/memory-proving-queue.js';
 import { ProverAgent } from '../prover-agent/prover-agent.js';
-import { getEnvironmentConfig, getSimulationProvider, makeGlobals } from './fixtures.js';
+import { getEnvironmentConfig, getSimulationProvider, makeGlobals, updateExpectedTreesFromTxs } from './fixtures.js';
 
 export class TestContext {
   private headers: Map<number, BlockHeader> = new Map();
@@ -78,7 +84,7 @@ export class TestContext {
 
     worldStateDB.getMerkleInterface.mockReturnValue(publicDb);
 
-    const publicTxSimulator = new PublicTxSimulator(publicDb, worldStateDB, telemetry, globalVariables);
+    const publicTxSimulator = new PublicTxSimulator(publicDb, worldStateDB, telemetry, globalVariables, true);
     const processor = new PublicProcessor(
       publicDb,
       globalVariables,
@@ -179,6 +185,7 @@ export class TestContext {
     const txs = times(numTxs, i =>
       this.makeProcessedTx({ seed: i + blockNum * 1000, globalVariables, ...makeProcessedTxOpts(i) }),
     );
+    await this.setEndTreeRoots(txs);
 
     const block = await buildBlock(txs, globalVariables, msgs, db);
     this.headers.set(blockNum, block.header);
@@ -216,6 +223,22 @@ export class TestContext {
     );
   }
 
+  public async setEndTreeRoots(txs: ProcessedTx[]) {
+    const db = await this.worldState.fork();
+    for (const tx of txs) {
+      await updateExpectedTreesFromTxs(db, [tx]);
+      const stateReference = await db.getStateReference();
+      if (tx.avmProvingRequest) {
+        tx.avmProvingRequest.inputs.output.endTreeSnapshots = new TreeSnapshots(
+          stateReference.l1ToL2MessageTree,
+          stateReference.partial.noteHashTree,
+          stateReference.partial.nullifierTree,
+          stateReference.partial.publicDataTree,
+        );
+      }
+    }
+  }
+
   private async processPublicFunctionsWithMockExecutorImplementation(
     txs: Tx[],
     maxTransactions: number,

yarn-project/prover-client/src/orchestrator/orchestrator_errors.test.ts

@@ -25,6 +25,7 @@ describe('prover/orchestrator/errors', () => {
   describe('errors', () => {
     it('throws if adding too many transactions', async () => {
       const txs = times(4, i => context.makeProcessedTx(i + 1));
+      await context.setEndTreeRoots(txs);
 
       orchestrator.startNewEpoch(1, 1, 1);
       await orchestrator.startNewBlock(context.globalVariables, []);