fix: avm gas and non-member (#10709)

This PR fixes two things:

1) Non-membership checks weren't properly being handled in the witgen
(the assertion was too strict)
2) End gas handling

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

@@ -78,6 +78,31 @@ FF AvmMerkleTreeTraceBuilder::unconstrained_update_leaf_index(const FF& leaf_val
     return unconstrained_compute_root_from_path(leaf_value, leaf_index, path);
 }
 
+bool AvmMerkleTreeTraceBuilder::assert_public_data_non_membership_check(
+    const PublicDataTreeLeafPreimage& low_leaf_preimage, const FF& leaf_slot)
+{
+    auto low_leaf_slot = uint256_t(low_leaf_preimage.slot);
+    auto low_leaf_next_index = uint256_t(low_leaf_preimage.next_index);
+    auto low_leaf_next_slot = uint256_t(low_leaf_preimage.next_slot);
+
+    auto leaf_slot_value = uint256_t(leaf_slot);
+
+    return low_leaf_slot < leaf_slot_value && (low_leaf_next_index == 0 || low_leaf_next_slot > leaf_slot_value);
+}
+
+bool AvmMerkleTreeTraceBuilder::assert_nullifier_non_membership_check(const NullifierLeafPreimage& low_leaf_preimage,
+                                                                      const FF& siloed_nullifier)
+{
+    auto low_leaf_nullifier = uint256_t(low_leaf_preimage.nullifier);
+    auto low_leaf_next_index = uint256_t(low_leaf_preimage.next_index);
+    auto low_leaf_next_nullifier = uint256_t(low_leaf_preimage.next_nullifier);
+
+    auto siloed_leaf_nullifier = uint256_t(siloed_nullifier);
+
+    return low_leaf_nullifier < siloed_leaf_nullifier &&
+           (low_leaf_next_index == 0 || low_leaf_next_nullifier > siloed_leaf_nullifier);
+}
+
 /**************************************************************************************************
  *                          STORAGE TREE OPERATIONS
  **************************************************************************************************/
@@ -113,6 +138,10 @@ FF AvmMerkleTreeTraceBuilder::perform_storage_write([[maybe_unused]] uint32_t cl
             unconstrained_update_leaf_index(low_preimage_hash, static_cast<uint64_t>(low_index), low_path);
         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
+    // value)
+    bool low_leaf_conditions = assert_public_data_non_membership_check(low_preimage, slot);
+    ASSERT(low_leaf_conditions);
     // The new leaf for an insertion is
     PublicDataTreeLeafPreimage new_preimage{
         .slot = slot, .value = value, .next_index = low_preimage.next_index, .next_slot = low_preimage.next_slot
@@ -165,6 +194,10 @@ FF AvmMerkleTreeTraceBuilder::perform_nullifier_append([[maybe_unused]] uint32_t
         ASSERT(is_member);
         return tree_snapshots.nullifier_tree.root;
     }
+    // Check the low leaf conditions (i.e. the slot is sandwiched by the low nullifier, or the new slot is a max
+    // value)
+    bool low_leaf_conditions = assert_nullifier_non_membership_check(low_preimage, nullifier);
+    ASSERT(low_leaf_conditions);
     // Check membership of the low leaf
     bool low_leaf_member = unconstrained_check_membership(
         low_preimage_hash, static_cast<uint64_t>(low_index), low_path, tree_snapshots.nullifier_tree.root);

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

@@ -101,6 +101,11 @@ class AvmMerkleTreeTraceBuilder {
     static FF unconstrained_silo_note_hash(FF contract_address, FF note_hash);
     static FF unconstrained_silo_nullifier(FF contract_address, FF nullifier);
     static FF unconstrained_compute_public_tree_leaf_slot(FF contract_address, FF leaf_index);
+    // Could probably template these
+    static bool assert_public_data_non_membership_check(const PublicDataTreeLeafPreimage& low_leaf_preimage,
+                                                        const FF& leaf_slot);
+    static bool assert_nullifier_non_membership_check(const NullifierLeafPreimage& low_leaf_preimage,
+                                                      const FF& nullifier);
 
     void finalize(std::vector<AvmFullRow<FF>>& main_trace);
 

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

@@ -2535,19 +2535,25 @@ AvmError AvmTraceBuilder::op_sload(uint8_t indirect, uint32_t slot_offset, uint3
 
     // Retrieve the public data read hint for this sload
     PublicDataReadTreeHint read_hint = execution_hints.storage_read_hints.at(storage_read_counter++);
+    // Check that the hinted leaf is a member of the public data tree
+    bool is_member = merkle_tree_trace_builder.perform_storage_read(
+        clk, read_hint.leaf_preimage, read_hint.leaf_index, read_hint.sibling_path);
+    ASSERT(is_member);
 
     // Compute the tree slot
     FF computed_tree_slot =
         merkle_tree_trace_builder.compute_public_tree_leaf_slot(clk, current_ext_call_ctx.contract_address, read_slot);
-    // Sanity check that the computed slot using the value read from slot_offset should match the read hint
-    ASSERT(computed_tree_slot == read_hint.leaf_preimage.slot);
-
-    // Check that the leaf is a member of the public data tree
-    bool is_member = merkle_tree_trace_builder.perform_storage_read(
-        clk, read_hint.leaf_preimage, read_hint.leaf_index, read_hint.sibling_path);
-    ASSERT(is_member);
+    // Check if the computed_tree_slot matches the read hint
+    bool exists = computed_tree_slot == read_hint.leaf_preimage.slot;
+
+    // If it doesnt exist, we should check the low nullifier conditions
+    if (!exists) {
+        bool non_member = AvmMerkleTreeTraceBuilder::assert_public_data_non_membership_check(read_hint.leaf_preimage,
+                                                                                             computed_tree_slot);
+        ASSERT(non_member);
+    }
 
-    FF value = read_hint.leaf_preimage.value;
+    FF value = exists ? read_hint.leaf_preimage.value : FF::zero();
     auto write_a = constrained_write_to_memory(
         call_ptr, clk, resolved_dest, value, AvmMemoryTag::FF, AvmMemoryTag::FF, IntermRegister::IA);
 
@@ -2836,7 +2842,6 @@ AvmError AvmTraceBuilder::op_nullifier_exists(uint8_t indirect,
     Row row;
 
     // Exists is written to b
-    bool exists = false;
     if (is_ok(error)) {
         NullifierReadTreeHint nullifier_read_hint = execution_hints.nullifier_read_hints.at(nullifier_read_counter++);
         FF nullifier_value = unconstrained_read_from_memory(resolved_nullifier_offset);
@@ -2847,17 +2852,11 @@ AvmError AvmTraceBuilder::op_nullifier_exists(uint8_t indirect,
                                                                           nullifier_read_hint.low_leaf_index,
                                                                           nullifier_read_hint.low_leaf_sibling_path);
         ASSERT(is_member);
-
-        if (siloed_nullifier == nullifier_read_hint.low_leaf_preimage.nullifier) {
-            // This is a direct membership check
-            exists = true;
-        } else {
-            exists = false;
-            // This is a non-membership proof
-            // Show that the target nullifier meets the non membership conditions (sandwich or max)
-            ASSERT(siloed_nullifier < nullifier_read_hint.low_leaf_preimage.nullifier &&
-                   (nullifier_read_hint.low_leaf_preimage.next_nullifier == FF::zero() ||
-                    siloed_nullifier > nullifier_read_hint.low_leaf_preimage.next_nullifier));
+        bool exists = siloed_nullifier == nullifier_read_hint.low_leaf_preimage.nullifier;
+        if (!exists) {
+            bool is_non_member = AvmMerkleTreeTraceBuilder::assert_nullifier_non_membership_check(
+                nullifier_read_hint.low_leaf_preimage, siloed_nullifier);
+            ASSERT(is_non_member);
         }
 
         auto read_a = constrained_read_from_memory(
@@ -4764,8 +4763,15 @@ std::vector<Row> AvmTraceBuilder::finalize(bool apply_end_gas_assertions)
 
     if (apply_end_gas_assertions) {
         // Sanity check that the amount of gas consumed matches what we expect from the public inputs
+        auto const l2_gas_left_after_private =
+            public_inputs.gas_settings.gas_limits.l2_gas - public_inputs.start_gas_used.l2_gas;
+        auto const allocated_l2_gas =
+            std::min(l2_gas_left_after_private, static_cast<uint32_t>(MAX_L2_GAS_PER_TX_PUBLIC_PORTION));
+        // Since end_gas_used also contains start_gas_used, we need to subtract it out to see what is consumed by the
+        // public computation only
+        auto expected_public_gas_consumed = public_inputs.end_gas_used.l2_gas - public_inputs.start_gas_used.l2_gas;
+        auto expected_end_gas_l2 = allocated_l2_gas - expected_public_gas_consumed;
         auto last_l2_gas_remaining = main_trace.back().main_l2_gas_remaining;
-        auto expected_end_gas_l2 = public_inputs.gas_settings.gas_limits.l2_gas - public_inputs.end_gas_used.l2_gas;
         vinfo("Last L2 gas remaining: ", last_l2_gas_remaining);
         vinfo("Expected end gas L2: ", expected_end_gas_l2);
         ASSERT(last_l2_gas_remaining == expected_end_gas_l2);