fix: small fixes to avm witgen (#10749)

- Add tree snapshots to nested call contexts
- Handle DA limits properly in witgen
- Respect staticcall in tree operations

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

@@ -21,6 +21,7 @@ enum class AvmError : uint32_t {
     DUPLICATE_NULLIFIER,
     SIDE_EFFECT_LIMIT_REACHED,
     OUT_OF_GAS,
+    STATIC_CALL_ALTERATION
 };
 
 } // namespace bb::avm_trace

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

@@ -146,6 +146,8 @@ FF AvmMerkleTreeTraceBuilder::perform_storage_write([[maybe_unused]] uint32_t cl
         // Update the low leaf
         tree_snapshots.public_data_tree.root =
             unconstrained_update_leaf_index(low_preimage_hash, static_cast<uint64_t>(low_index), low_path);
+        // Update the set of writes to unique slots
+        public_data_unique_writes.insert(slot);
         return tree_snapshots.public_data_tree.root;
     }
     // Check the low leaf conditions (i.e. the slot is sandwiched by the low nullifier, or the new slot is a max
@@ -172,6 +174,8 @@ FF AvmMerkleTreeTraceBuilder::perform_storage_write([[maybe_unused]] uint32_t cl
     // Insert the new leaf into the tree
     tree_snapshots.public_data_tree.root = unconstrained_update_leaf_index(leaf_preimage_hash, index, insertion_path);
     tree_snapshots.public_data_tree.size++;
+    // Update the set of writes to unique slots
+    public_data_unique_writes.insert(slot);
     return tree_snapshots.public_data_tree.root;
 }
 

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

@@ -6,6 +6,7 @@
 #include "barretenberg/vm/avm/trace/gadgets/poseidon2.hpp"
 #include "barretenberg/vm/avm/trace/public_inputs.hpp"
 
+#include <unordered_set>
 #include <vector>
 namespace bb::avm_trace {
 
@@ -42,6 +43,12 @@ class AvmMerkleTreeTraceBuilder {
     FF compute_public_tree_leaf_slot(uint32_t clk, FF contract_address, FF leaf_index);
 
     TreeSnapshots& get_tree_snapshots() { return tree_snapshots; }
+    void restore_tree_state(TreeSnapshots& tree_snapshots, std::unordered_set<FF>& public_data_unique_writes)
+    {
+        this->tree_snapshots = tree_snapshots;
+        this->public_data_unique_writes = public_data_unique_writes;
+    }
+    std::unordered_set<FF> get_public_data_unique_writes() { return public_data_unique_writes; }
 
     // Public Data Tree
     bool perform_storage_read(uint32_t clk,
@@ -120,6 +127,7 @@ class AvmMerkleTreeTraceBuilder {
     std::vector<MerkleEntry> merkle_check_trace;
     TreeSnapshots non_revertible_tree_snapshots;
     TreeSnapshots tree_snapshots;
+    std::unordered_set<FF> public_data_unique_writes;
     MerkleEntry compute_root_from_path(uint32_t clk,
                                        const FF& leaf_value,
                                        const uint64_t leaf_index,

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

@@ -100,7 +100,9 @@ uint32_t finalize_rng_chks_for_testing(std::vector<Row>& main_trace,
     auto old_size = main_trace.size();
     for (auto const& clk : custom_clk) {
         if (clk >= old_size) {
-            main_trace.push_back(Row{ .main_clk = FF(clk) });
+            Row row{};
+            row.main_clk = clk;
+            main_trace.push_back(row);
         }
     }
 
@@ -219,6 +221,19 @@ uint32_t AvmTraceBuilder::get_inserted_note_hashes_count()
            public_inputs.start_tree_snapshots.note_hash_tree.size;
 }
 
+uint32_t AvmTraceBuilder::get_inserted_nullifiers_count()
+{
+    return merkle_tree_trace_builder.get_tree_snapshots().nullifier_tree.size -
+           public_inputs.start_tree_snapshots.nullifier_tree.size;
+}
+
+// Keeping track of the public data writes isn't as simple as comparing sizes
+// because of updates to leaves that were already in the tree and state squashing
+uint32_t AvmTraceBuilder::get_public_data_writes_count()
+{
+    return static_cast<uint32_t>(merkle_tree_trace_builder.get_public_data_unique_writes().size());
+}
+
 void AvmTraceBuilder::insert_private_state(const std::vector<FF>& siloed_nullifiers,
                                            const std::vector<FF>& unique_note_hashes)
 {
@@ -305,6 +320,7 @@ void AvmTraceBuilder::pay_fee()
         throw std::runtime_error("Not enough balance for fee payer to pay for transaction");
     }
 
+    // TS equivalent:
     // writeStorage(FEE_JUICE_ADDRESS, balance_slot, updated_balance);
     PublicDataWriteTreeHint write_hint = execution_hints.storage_write_hints.at(storage_write_counter++);
     ASSERT(write_hint.new_leaf_preimage.value == updated_balance);
@@ -2818,15 +2834,17 @@ AvmError AvmTraceBuilder::op_sload(uint8_t indirect, uint32_t slot_offset, uint3
 
 AvmError AvmTraceBuilder::op_sstore(uint8_t indirect, uint32_t src_offset, uint32_t slot_offset)
 {
-    // We keep the first encountered error
-    AvmError error = AvmError::NO_ERROR;
     auto clk = static_cast<uint32_t>(main_trace.size()) + 1;
+    uint32_t unique_public_data_slot_writes = get_public_data_writes_count();
+    AvmError error = current_ext_call_ctx.is_static_call ? AvmError::STATIC_CALL_ALTERATION
+                     : unique_public_data_slot_writes >= MAX_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX
+                         ? AvmError::SIDE_EFFECT_LIMIT_REACHED
+                         : AvmError::NO_ERROR;
 
-    if (storage_write_counter >= MAX_PUBLIC_DATA_UPDATE_REQUESTS_PER_TX) {
+    if (!is_ok(error)) {
         // NOTE: the circuit constraint for this limit should only be applied
         // for the storage writes performed by this opcode. An exception should before
         // made for the fee juice storage write made after teardown.
-        error = AvmError::SIDE_EFFECT_LIMIT_REACHED;
         auto row = Row{
             .main_clk = clk,
             .main_internal_return_ptr = internal_return_ptr,
@@ -3005,8 +3023,11 @@ AvmError AvmTraceBuilder::op_emit_note_hash(uint8_t indirect, uint32_t note_hash
 {
     auto const clk = static_cast<uint32_t>(main_trace.size()) + 1;
     uint32_t inserted_note_hashes_count = get_inserted_note_hashes_count();
-    if (inserted_note_hashes_count >= MAX_NOTE_HASHES_PER_TX) {
-        AvmError error = AvmError::SIDE_EFFECT_LIMIT_REACHED;
+    AvmError error = current_ext_call_ctx.is_static_call                    ? AvmError::STATIC_CALL_ALTERATION
+                     : inserted_note_hashes_count >= MAX_NOTE_HASHES_PER_TX ? AvmError::SIDE_EFFECT_LIMIT_REACHED
+                                                                            : AvmError::NO_ERROR;
+
+    if (!is_ok(error)) {
         auto row = Row{
             .main_clk = clk,
             .main_internal_return_ptr = internal_return_ptr,
@@ -3019,8 +3040,12 @@ AvmError AvmTraceBuilder::op_emit_note_hash(uint8_t indirect, uint32_t note_hash
         return error;
     }
 
-    auto [row, error] = create_kernel_output_opcode(indirect, clk, note_hash_offset);
+    auto [row, output_error] = create_kernel_output_opcode(indirect, clk, note_hash_offset);
     row.main_sel_op_emit_note_hash = FF(1);
+    if (is_ok(error)) {
+        error = output_error;
+    }
+
     row.main_op_err = FF(static_cast<uint32_t>(!is_ok(error)));
 
     AppendTreeHint note_hash_write_hint = execution_hints.note_hash_write_hints.at(note_hash_write_counter++);
@@ -3157,12 +3182,14 @@ AvmError AvmTraceBuilder::op_nullifier_exists(uint8_t indirect,
 
 AvmError AvmTraceBuilder::op_emit_nullifier(uint8_t indirect, uint32_t nullifier_offset)
 {
-    // We keep the first encountered error
-    AvmError error = AvmError::NO_ERROR;
     auto const clk = static_cast<uint32_t>(main_trace.size()) + 1;
 
-    if (nullifier_write_counter >= MAX_NULLIFIERS_PER_TX) {
-        error = AvmError::SIDE_EFFECT_LIMIT_REACHED;
+    uint32_t inserted_nullifier_count = get_inserted_nullifiers_count();
+    AvmError error = current_ext_call_ctx.is_static_call                 ? AvmError::STATIC_CALL_ALTERATION
+                     : inserted_nullifier_count >= MAX_NULLIFIERS_PER_TX ? AvmError::SIDE_EFFECT_LIMIT_REACHED
+                                                                         : AvmError::NO_ERROR;
+
+    if (!is_ok(error)) {
         auto row = Row{
             .main_clk = clk,
             .main_internal_return_ptr = internal_return_ptr,
@@ -3754,7 +3781,9 @@ AvmError AvmTraceBuilder::constrain_external_call(OpCode opcode,
 
     // We push the current ext call ctx onto the stack and initialize a new one
     current_ext_call_ctx.last_pc = pc;
-    current_ext_call_ctx.success_offset = resolved_success_offset,
+    current_ext_call_ctx.success_offset = resolved_success_offset;
+    current_ext_call_ctx.tree_snapshot = merkle_tree_trace_builder.get_tree_snapshots();
+    current_ext_call_ctx.public_data_unique_writes = merkle_tree_trace_builder.get_public_data_unique_writes();
     external_call_ctx_stack.emplace(current_ext_call_ctx);
 
     // Ext Ctx setup
@@ -3771,6 +3800,7 @@ AvmError AvmTraceBuilder::constrain_external_call(OpCode opcode,
     current_ext_call_ctx = ExtCallCtx{
         .context_id = static_cast<uint8_t>(clk),
         .parent_id = current_ext_call_ctx.context_id,
+        .is_static_call = opcode == OpCode::STATICCALL,
         .contract_address = read_addr.val,
         .calldata = calldata,
         .nested_returndata = {},
@@ -3781,6 +3811,7 @@ AvmError AvmTraceBuilder::constrain_external_call(OpCode opcode,
         .l2_gas_left = l2_gas_allocated_to_nested_call,
         .da_gas_left = da_gas_allocated_to_nested_call,
         .internal_return_ptr_stack = {},
+        .tree_snapshot = {},
     };
 
     allocate_gas_for_call(l2_gas_allocated_to_nested_call, da_gas_allocated_to_nested_call);
@@ -4079,6 +4110,8 @@ ReturnDataError AvmTraceBuilder::op_revert(uint8_t indirect, uint32_t ret_offset
             set_call_ptr(static_cast<uint8_t>(current_ext_call_ctx.context_id));
             write_to_memory(
                 current_ext_call_ctx.success_offset, /*success=*/FF::one(), AvmMemoryTag::U1, /*fix_pc=*/false);
+            merkle_tree_trace_builder.restore_tree_state(current_ext_call_ctx.tree_snapshot,
+                                                         current_ext_call_ctx.public_data_unique_writes);
         }
     }
 
@@ -4091,7 +4124,7 @@ ReturnDataError AvmTraceBuilder::op_revert(uint8_t indirect, uint32_t ret_offset
             .main_internal_return_ptr = FF(internal_return_ptr),
             .main_op_err = FF(static_cast<uint32_t>(!is_ok(error))),
             .main_pc = pc,
-            .main_sel_op_external_return = 1,
+            .main_sel_op_external_revert = 1,
         });
 
         pc = is_top_level ? UINT32_MAX : current_ext_call_ctx.last_pc;

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

@@ -266,6 +266,7 @@ class AvmTraceBuilder {
     struct ExtCallCtx {
         uint32_t context_id; // This is the unique id of the ctx, we'll use the clk
         uint32_t parent_id;
+        bool is_static_call = false;
         FF contract_address{};
         std::vector<FF> calldata;
         std::vector<FF> nested_returndata;
@@ -276,6 +277,8 @@ class AvmTraceBuilder {
         uint32_t l2_gas_left; // as of start of latest nested call
         uint32_t da_gas_left; // as of start of latest nested call
         std::stack<uint32_t> internal_return_ptr_stack;
+        TreeSnapshots tree_snapshot; // This is the tree state at the time of the call
+        std::unordered_set<FF> public_data_unique_writes;
     };
 
     ExtCallCtx current_ext_call_ctx{};
@@ -367,6 +370,8 @@ class AvmTraceBuilder {
                                       IntermRegister reg,
                                       AvmMemTraceBuilder::MemOpOwner mem_op_owner = AvmMemTraceBuilder::MAIN);
     uint32_t get_inserted_note_hashes_count();
+    uint32_t get_inserted_nullifiers_count();
+    uint32_t get_public_data_writes_count();
     FF get_tx_hash() const { return public_inputs.previous_non_revertible_accumulated_data.nullifiers[0]; }
 
     // TODO: remove these once everything is constrained.