From a56d282b8ff41bdd71c10da3fe6da150c80a15c3 Mon Sep 17 00:00:00 2001
From: dbanks12 <david@aztec-labs.com>
Date: Wed, 6 Aug 2025 20:38:35 +0000
Subject: [PATCH] fix!: some missing context stuff in the AVM Cpp

---
 .../cpp/src/barretenberg/vm2/simulation/execution.cpp      | 2 +-
 .../cpp/src/barretenberg/vm2/simulation/execution.test.cpp | 1 -
 .../cpp/src/barretenberg/vm2/tracegen/execution_trace.cpp  | 5 +++++
 barretenberg/cpp/src/barretenberg/vm2/tracegen_helper.cpp  | 7 +++++++
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.cpp b/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.cpp
index 4c67319e5d41..8debbcd93266 100644
--- a/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.cpp
+++ b/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.cpp
@@ -1030,7 +1030,7 @@ void Execution::handle_enter_call(ContextInterface& parent_context, std::unique_
     ctx_stack_events.emit(
         { .id = parent_context.get_context_id(),
           .parent_id = parent_context.get_parent_id(),
-          .entered_context_id = context_provider.get_next_context_id(),
+          .entered_context_id = child_context->get_context_id(), // gets the context id of the child!
           .next_pc = parent_context.get_next_pc(),
           .msg_sender = parent_context.get_msg_sender(),
           .contract_addr = parent_context.get_address(),
diff --git a/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.test.cpp b/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.test.cpp
index 8eaac4444fe1..0fb3876a1ab3 100644
--- a/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.test.cpp
+++ b/barretenberg/cpp/src/barretenberg/vm2/simulation/execution.test.cpp
@@ -205,7 +205,6 @@ TEST_F(ExecutionSimulationTest, Call)
 
     // Context snapshotting
     EXPECT_CALL(context, get_context_id);
-    EXPECT_CALL(context_provider, get_next_context_id);
     EXPECT_CALL(context, get_parent_id);
     EXPECT_CALL(context, get_next_pc);
     EXPECT_CALL(context, get_is_static);
diff --git a/barretenberg/cpp/src/barretenberg/vm2/tracegen/execution_trace.cpp b/barretenberg/cpp/src/barretenberg/vm2/tracegen/execution_trace.cpp
index 1f613a556110..e62e4c3d2a9b 100644
--- a/barretenberg/cpp/src/barretenberg/vm2/tracegen/execution_trace.cpp
+++ b/barretenberg/cpp/src/barretenberg/vm2/tracegen/execution_trace.cpp
@@ -20,6 +20,7 @@
 #include "barretenberg/vm2/generated/columns.hpp"
 #include "barretenberg/vm2/generated/relations/lookups_addressing.hpp"
 #include "barretenberg/vm2/generated/relations/lookups_alu.hpp"
+#include "barretenberg/vm2/generated/relations/lookups_context.hpp"
 #include "barretenberg/vm2/generated/relations/lookups_emit_notehash.hpp"
 #include "barretenberg/vm2/generated/relations/lookups_emit_nullifier.hpp"
 #include "barretenberg/vm2/generated/relations/lookups_execution.hpp"
@@ -1177,6 +1178,10 @@ const InteractionDefinition ExecutionTraceBuilder::interactions =
         .add<lookup_execution_check_radix_gt_256_settings, InteractionType::LookupGeneric>()
         .add<lookup_execution_get_p_limbs_settings, InteractionType::LookupGeneric>()
         .add<lookup_execution_get_max_limbs_settings, InteractionType::LookupGeneric>()
+        // Context Stack
+        .add<lookup_context_ctx_stack_call_settings, InteractionType::LookupGeneric>()
+        .add<lookup_context_ctx_stack_rollback_settings, InteractionType::LookupGeneric>()
+        .add<lookup_context_ctx_stack_return_settings, InteractionType::LookupGeneric>()
         // External Call
         .add<lookup_external_call_call_allocated_left_l2_range_settings, InteractionType::LookupIntoIndexedByClk>()
         .add<lookup_external_call_call_allocated_left_da_range_settings, InteractionType::LookupIntoIndexedByClk>()
diff --git a/barretenberg/cpp/src/barretenberg/vm2/tracegen_helper.cpp b/barretenberg/cpp/src/barretenberg/vm2/tracegen_helper.cpp
index ec07a10275c4..a8d724b3fa09 100644
--- a/barretenberg/cpp/src/barretenberg/vm2/tracegen_helper.cpp
+++ b/barretenberg/cpp/src/barretenberg/vm2/tracegen_helper.cpp
@@ -20,6 +20,7 @@
 #include "barretenberg/vm2/tracegen/bytecode_trace.hpp"
 #include "barretenberg/vm2/tracegen/calldata_trace.hpp"
 #include "barretenberg/vm2/tracegen/class_id_derivation_trace.hpp"
+#include "barretenberg/vm2/tracegen/context_stack_trace.hpp"
 #include "barretenberg/vm2/tracegen/contract_instance_retrieval_trace.hpp"
 #include "barretenberg/vm2/tracegen/data_copy_trace.hpp"
 #include "barretenberg/vm2/tracegen/ecc_trace.hpp"
@@ -388,6 +389,12 @@ void AvmTraceGenHelper::fill_trace_columns(TraceContainer& trace,
                                    internal_call_stack_builder.process(events.internal_call_stack_events, trace));
                     clear_events(events.internal_call_stack_events);
                 },
+                [&]() {
+                    ContextStackTraceBuilder context_stack_builder;
+                    AVM_TRACK_TIME("tracegen/context_stack",
+                                   context_stack_builder.process(events.context_stack, trace));
+                    clear_events(events.context_stack);
+                },
                 [&]() {
                     NoteHashTreeCheckTraceBuilder note_hash_tree_check_trace_builder;
                     AVM_TRACK_TIME(