feat: Add tree equality assertions

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

@@ -182,6 +182,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,

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

@@ -1,5 +1,5 @@
 import { type IndexedTreeId, MerkleTreeId, type MerkleTreeReadOperations, getTreeHeight } from '@aztec/circuit-types';
-import { NullifierLeafPreimage, PublicDataTreeLeafPreimage } from '@aztec/circuits.js';
+import { AppendOnlyTreeSnapshot, NullifierLeafPreimage, PublicDataTreeLeafPreimage } from '@aztec/circuits.js';
 import { poseidon2Hash } from '@aztec/foundation/crypto';
 import { Fr } from '@aztec/foundation/fields';
 import { type IndexedTreeLeafPreimage, type TreeLeafPreimage } from '@aztec/foundation/trees';
@@ -550,6 +550,11 @@ export class AvmEphemeralForest {
     const input = preimage.toHashInputs().map(x => Fr.fromBuffer(x));
     return poseidon2Hash(input);
   }
+
+  getTreeSnapshot(id: MerkleTreeId): AppendOnlyTreeSnapshot {
+    const noteHashTree = this.treeMap.get(id)!;
+    return new AppendOnlyTreeSnapshot(noteHashTree.getRoot(), Number(noteHashTree.leafCount));
+  }
 }
 
 /****************************************************/

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

@@ -10,7 +10,6 @@ import {
   TxHash,
 } from '@aztec/circuit-types';
 import {
-  AppendOnlyTreeSnapshot,
   AvmCircuitInputs,
   type AvmCircuitPublicInputs,
   type AztecAddress,
@@ -19,6 +18,8 @@ import {
   type GasSettings,
   type GlobalVariables,
   MAX_L2_GAS_PER_TX_PUBLIC_PORTION,
+  MAX_NOTE_HASHES_PER_TX,
+  MAX_NULLIFIERS_PER_TX,
   type PrivateToPublicAccumulatedData,
   type PublicCallRequest,
   PublicCircuitPublicInputs,
@@ -311,16 +312,16 @@ export class PublicTxContext {
    */
   private generateAvmCircuitPublicInputs(endStateReference: StateReference): AvmCircuitPublicInputs {
     assert(this.halted, 'Can only get AvmCircuitPublicInputs after tx execution ends');
-    const ephemeralTrees = this.state.getActiveStateManager().merkleTrees.treeMap;
-
-    const getAppendSnaphot = (id: MerkleTreeId) => {
-      const tree = ephemeralTrees.get(id)!;
-      return new AppendOnlyTreeSnapshot(tree.getRoot(), Number(tree.leafCount));
-    };
-
-    const noteHashTree = getAppendSnaphot(MerkleTreeId.NOTE_HASH_TREE);
-    const nullifierTree = getAppendSnaphot(MerkleTreeId.NULLIFIER_TREE);
-    const publicDataTree = getAppendSnaphot(MerkleTreeId.PUBLIC_DATA_TREE);
+    const ephemeralTrees = this.state.getActiveStateManager().merkleTrees;
+
+    const noteHashTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.NOTE_HASH_TREE);
+    const nullifierTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.NULLIFIER_TREE);
+    const publicDataTree = ephemeralTrees.getTreeSnapshot(MerkleTreeId.PUBLIC_DATA_TREE);
+    // Pad the note hash and nullifier trees
+    noteHashTree.nextAvailableLeafIndex =
+      this.startStateReference.partial.noteHashTree.nextAvailableLeafIndex + MAX_NOTE_HASHES_PER_TX;
+    nullifierTree.nextAvailableLeafIndex =
+      this.startStateReference.partial.nullifierTree.nextAvailableLeafIndex + MAX_NULLIFIERS_PER_TX;
 
     const endTreeSnapshots = new TreeSnapshots(
       endStateReference.l1ToL2MessageTree,