fix: resolve misc bugs handling phases in avm witgen

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

@@ -403,6 +403,7 @@ std::vector<Row> Execution::gen_trace(AvmPublicInputs const& public_inputs,
 
     // Loop over all the public call requests
     auto const phases = { TxExecutionPhase::SETUP, TxExecutionPhase::APP_LOGIC, TxExecutionPhase::TEARDOWN };
+    size_t enqueued_call_hint_index = 0;
     for (auto phase : phases) {
         const auto public_call_requests = phase == TxExecutionPhase::SETUP       ? setup_call_requests
                                           : phase == TxExecutionPhase::APP_LOGIC ? app_logic_call_requests
@@ -447,11 +448,11 @@ std::vector<Row> Execution::gen_trace(AvmPublicInputs const& public_inputs,
             auto public_call_request = public_call_requests.at(i);
             trace_builder.set_public_call_request(public_call_request);
             // At the start of each enqueued call, we read the enqueued call hints
-            auto enqueued_call_hint = execution_hints.enqueued_call_hints.at(i);
+            auto enqueued_call_hint = execution_hints.enqueued_call_hints.at(enqueued_call_hint_index++);
             ASSERT(public_call_request.contract_address == enqueued_call_hint.contract_address);
             // Execute!
-            phase_error =
-                Execution::execute_enqueued_call(trace_builder, enqueued_call_hint, returndata, apply_e2e_assertions);
+            phase_error = Execution::execute_enqueued_call(
+                phase, trace_builder, enqueued_call_hint, returndata, apply_e2e_assertions);
 
             if (!is_ok(phase_error)) {
                 info("Phase ", to_name(phase), " reverted.");
@@ -468,6 +469,13 @@ std::vector<Row> Execution::gen_trace(AvmPublicInputs const& public_inputs,
             throw std::runtime_error("A revert was encountered in the SETUP phase, killing the entire TX");
             break;
         }
+        vinfo("Ended phase ",
+              to_name(phase),
+              " with ",
+              trace_builder.get_l2_gas_left(),
+              " L2 gas left and ",
+              trace_builder.get_da_gas_left(),
+              " DA gas left");
     }
 
     if (apply_e2e_assertions) {
@@ -493,13 +501,18 @@ std::vector<Row> Execution::gen_trace(AvmPublicInputs const& public_inputs,
  * @returns the error/result of the enqueued call
  *
  */
-AvmError Execution::execute_enqueued_call(AvmTraceBuilder& trace_builder,
+AvmError Execution::execute_enqueued_call(TxExecutionPhase phase,
+                                          AvmTraceBuilder& trace_builder,
                                           AvmEnqueuedCallHint& enqueued_call_hint,
                                           std::vector<FF>& returndata,
                                           bool check_bytecode_membership)
 {
     AvmError error = AvmError::NO_ERROR;
 
+    // save gas before phase for use after teardown (teardown shouldn't  affect end gas)
+    const auto l2_gas_left_before_enqueued_call = trace_builder.get_l2_gas_left();
+    const auto da_gas_left_before_enqueued_call = trace_builder.get_da_gas_left();
+
     // These hints help us to set up first call ctx
     auto context_id = trace_builder.next_context_id;
     uint32_t l2_gas_allocated_to_enqueued_call = trace_builder.get_l2_gas_left();
@@ -533,7 +546,11 @@ AvmError Execution::execute_enqueued_call(AvmTraceBuilder& trace_builder,
         // This error was encountered before any opcodes were executed, but
         // we need at least one row in the execution trace to then mutate and say "it halted and consumed all gas!"
         trace_builder.op_add(0, 0, 0, 0, OpCode::ADD_8);
-        trace_builder.handle_exceptional_halt();
+        if (phase == TxExecutionPhase::TEARDOWN) {
+            trace_builder.handle_end_of_teardown(l2_gas_left_before_enqueued_call, da_gas_left_before_enqueued_call);
+        } else {
+            trace_builder.handle_exceptional_halt();
+        }
         return AvmError::FAILED_BYTECODE_RETRIEVAL;
     }
 
@@ -1124,7 +1141,12 @@ AvmError Execution::execute_enqueued_call(AvmTraceBuilder& trace_builder,
                  " IC: ",
                  error_ic);
 
-            trace_builder.handle_exceptional_halt();
+            // For nested calls or non-teardown-enqueued-calls, handle the exceptional halt
+            // (consume all gas). Don't  do it for teardown enqueued call because gas needs
+            // to be reset after teardown (teardown doesn't count towards end-gas).
+            if (!is_top_level || phase != TxExecutionPhase::TEARDOWN) {
+                trace_builder.handle_exceptional_halt();
+            }
 
             if (is_top_level) {
                 break;
@@ -1140,6 +1162,9 @@ AvmError Execution::execute_enqueued_call(AvmTraceBuilder& trace_builder,
             error = AvmError::NO_ERROR;
         }
     }
+    if (phase == TxExecutionPhase::TEARDOWN) {
+        trace_builder.handle_end_of_teardown(l2_gas_left_before_enqueued_call, da_gas_left_before_enqueued_call);
+    }
     return error;
 }
 

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

@@ -40,7 +40,8 @@ class Execution {
                                       ExecutionHints const& execution_hints,
                                       bool apply_e2e_assertions = false);
 
-    static AvmError execute_enqueued_call(AvmTraceBuilder& trace_builder,
+    static AvmError execute_enqueued_call(TxExecutionPhase phase,
+                                          AvmTraceBuilder& trace_builder,
                                           AvmEnqueuedCallHint& enqueued_call_hint,
                                           std::vector<FF>& returndata,
                                           bool check_bytecode_membership);

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

@@ -124,8 +124,8 @@ void AvmGasTraceBuilder::constrain_gas_for_halt(bool exceptional_halt,
                                                 uint32_t l2_gas_allocated_to_nested_call,
                                                 uint32_t da_gas_allocated_to_nested_call)
 {
-    debug("Resetting to parent's L2 gas left (", parent_l2_gas_left, ") before consuming gas allocated to nested call");
-    debug("Resetting to parent's DA gas left (", parent_da_gas_left, ") before consuming gas allocated to nested call");
+    debug("Resetting to parent's L2 gas left (", parent_l2_gas_left, ") before consuming gas used by nested call");
+    debug("Resetting to parent's DA gas left (", parent_da_gas_left, ") before consuming gas used by nested call");
     // how much gas did the nested call consume
     auto l2_gas_consumed = l2_gas_allocated_to_nested_call - remaining_l2_gas;
     auto da_gas_consumed = da_gas_allocated_to_nested_call - remaining_da_gas;

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

@@ -429,6 +429,22 @@ void AvmTraceBuilder::handle_exceptional_halt()
     }
 }
 
+void AvmTraceBuilder::handle_end_of_teardown(uint32_t pre_teardown_l2_gas_left, uint32_t pre_teardown_da_gas_left)
+{
+    vinfo("Handling end of teardown");
+
+    // modify the last row of the gas trace to reset back to pre-teardown gas
+    // since gas used by teardown doesn't contribute to end-gas
+    gas_trace_builder.constrain_gas_for_halt(/*exceptional_halt=*/true, // not really an exceptional halt
+                                             pre_teardown_l2_gas_left,
+                                             pre_teardown_da_gas_left,
+                                             /*l2_gas_allocated_to_nested_call=*/0,
+                                             /*da_gas_allocated_to_nested_call=*/0);
+
+    // max out the pc to signify "done"
+    pc = UINT32_MAX;
+}
+
 /**
  * @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

@@ -243,6 +243,7 @@ class AvmTraceBuilder {
     void pad_trees();
     void allocate_gas_for_call(uint32_t l2_gas, uint32_t da_gas);
     void handle_exceptional_halt();
+    void handle_end_of_teardown(uint32_t pre_teardown_l2_gas_left, uint32_t pre_teardown_da_gas_left);
 
     // These are used for testing only.
     AvmTraceBuilder& set_range_check_required(bool required)

yarn-project/bb-prover/src/avm_proving.test.ts

@@ -261,6 +261,22 @@ describe('AVM WitGen, proof generation and verification', () => {
     },
     TIMEOUT,
   );
+  it(
+    'Should prove and verify a TX that reverts in teardown',
+    async () => {
+      await proveAndVerifyAvmTestContract(
+        /*checkCircuitOnly=*/ true,
+        /*setupFunctionNames=*/ [],
+        /*setupArgs=*/ [],
+        /*appFunctionNames=*/ [],
+        /*appArgs=*/ [],
+        /*teardownFunctionName=*/ 'read_assert_storage_single',
+        /*teardownArgs=*/ [new Fr(10)],
+        /*expectRevert=*/ true,
+      );
+    },
+    TIMEOUT,
+  );
 });
 
 /**

yarn-project/simulator/src/public/fixtures/index.ts

@@ -79,16 +79,14 @@ export async function simulateAvmTestContractGenerateCircuitInputs(
     /*isStaticCall=*/ false,
   );
   const setupExecutionRequests: PublicExecutionRequest[] = [];
-  // we reverse order because the simulator expects it to be like a "stack" of calls to pop from
-  for (let i = setupFunctionNames.length - 1; i >= 0; i--) {
+  for (let i = 0; i < setupFunctionNames.length; i++) {
     const functionSelector = getAvmTestContractFunctionSelector(setupFunctionNames[i]);
     const fnArgs = [functionSelector.toField(), ...setupArgs[i]];
     const executionRequest = new PublicExecutionRequest(callContext, fnArgs);
     setupExecutionRequests.push(executionRequest);
   }
   const appExecutionRequests: PublicExecutionRequest[] = [];
-  // we reverse order because the simulator expects it to be like a "stack" of calls to pop from
-  for (let i = appFunctionNames.length - 1; i >= 0; i--) {
+  for (let i = 0; i < appFunctionNames.length; i++) {
     const functionSelector = getAvmTestContractFunctionSelector(appFunctionNames[i]);
     const fnArgs = [functionSelector.toField(), ...appArgs[i]];
     const executionRequest = new PublicExecutionRequest(callContext, fnArgs);
@@ -183,10 +181,11 @@ export function createTxForPublicCalls(
   // TODO(#9269): Remove this fake nullifier method as we move away from 1st nullifier as hash.
   forPublic.nonRevertibleAccumulatedData.nullifiers[0] = Fr.random(); // fake tx nullifier
 
-  for (let i = 0; i < setupExecutionRequests.length; i++) {
+  // We reverse order because the simulator expects it to be like a "stack" of calls to pop from
+  for (let i = setupCallRequests.length - 1; i >= 0; i--) {
     forPublic.nonRevertibleAccumulatedData.publicCallRequests[i] = setupCallRequests[i];
   }
-  for (let i = 0; i < appCallRequests.length; i++) {
+  for (let i = appCallRequests.length - 1; i >= 0; i--) {
     forPublic.revertibleAccumulatedData.publicCallRequests[i] = appCallRequests[i];
   }
   if (teardownExecutionRequest) {
@@ -207,12 +206,14 @@ export function createTxForPublicCalls(
   );
   const tx = Tx.newWithTxData(txData, teardownExecutionRequest);
 
-  for (let i = 0; i < setupExecutionRequests.length; i++) {
-    tx.enqueuedPublicFunctionCalls.push(setupExecutionRequests[i]);
-  }
-  for (let i = 0; i < appExecutionRequests.length; i++) {
+  // Reverse order because the simulator expects it to be like a "stack" of calls to pop from.
+  // Also push app calls before setup calls for this reason.
+  for (let i = appExecutionRequests.length - 1; i >= 0; i--) {
     tx.enqueuedPublicFunctionCalls.push(appExecutionRequests[i]);
   }
+  for (let i = setupExecutionRequests.length - 1; i >= 0; i--) {
+    tx.enqueuedPublicFunctionCalls.push(setupExecutionRequests[i]);
+  }
 
   return tx;
 }