feat: preserve loaded package debug info

Author: huitseeker

crates/lib/core/src/lib.rs

@@ -111,6 +111,7 @@ impl From<&CoreLibrary> for HostLibrary {
     fn from(core_lib: &CoreLibrary) -> Self {
         Self {
             mast_forest: core_lib.mast_forest().clone(),
+            package_debug_info: Ok(None),
             handlers: core_lib.handlers(),
         }
     }

crates/lib/core/tests/debug.rs

@@ -69,6 +69,7 @@ fn run(source: &str, advice: AdviceInputs) -> (String, ExecutionOutput) {
     let buf = Arc::new(Mutex::new(String::new()));
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers: debug_handlers_with_writer(SharedBuf(buf.clone())),
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");
@@ -293,6 +294,7 @@ fn print_mem_rejects_out_of_bounds_range_end() {
         .unwrap_program();
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers: debug_handlers_with_writer(SharedBuf(Arc::new(Mutex::new(String::new())))),
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");
@@ -335,6 +337,7 @@ fn print_mem_rejects_oversized_range() {
         .unwrap_program();
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers: debug_handlers_with_writer(SharedBuf(Arc::new(Mutex::new(String::new())))),
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");
@@ -377,6 +380,7 @@ fn print_mem_rejects_full_range() {
         .unwrap_program();
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers: debug_handlers_with_writer(SharedBuf(Arc::new(Mutex::new(String::new())))),
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");
@@ -580,6 +584,7 @@ fn debug_handlers_compose_with_default_core_handlers() {
     handlers.extend(advice_debug_handlers());
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers,
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");
@@ -649,6 +654,7 @@ fn noop_debug_handlers_run_print_stack_without_output() {
         .unwrap_program();
     let host_lib = HostLibrary {
         mast_forest: core_lib.mast_forest().clone(),
+        package_debug_info: Ok(None),
         handlers: noop_debug_handlers(),
     };
     let mut host = DefaultHost::default().with_library(host_lib).expect("failed to load host lib");

processor/src/continuation_stack.rs

@@ -1,7 +1,7 @@
-use alloc::vec::Vec;
+use alloc::{sync::Arc, vec::Vec};
 
 use miden_core::{mast::MastNodeId, program::Program};
-use miden_mast_package::debug_info::DebugSourceNodeId;
+use miden_mast_package::debug_info::{DebugSourceNodeId, PackageDebugInfo};
 
 /// A hint for the initial size of the continuation stack.
 const CONTINUATION_STACK_SIZE_HINT: usize = 64;
@@ -55,7 +55,10 @@ pub enum Continuation<F> {
     ///
     /// When we encounter an `ExternalNode`, we enter the corresponding MAST forest directly, and
     /// push an `EnterForest` continuation to restore the previous forest when done.
-    EnterForest(F),
+    EnterForest {
+        forest: F,
+        package_debug_info: Option<Arc<PackageDebugInfo>>,
+    },
 }
 
 impl<F> Continuation<F> {
@@ -82,7 +85,23 @@ impl<F> Continuation<F> {
             | Respan { node_id: _, batch_index: _ }
             | FinishBasicBlock(_) => true,
 
-            EnterForest(_) => false,
+            EnterForest { .. } => false,
+        }
+    }
+
+    #[cfg(any(test, feature = "testing"))]
+    pub(crate) fn exec_node(&self) -> Option<MastNodeId> {
+        match self {
+            Self::StartNode(node_id)
+            | Self::FinishJoin(node_id)
+            | Self::FinishSplit(node_id)
+            | Self::FinishLoop(node_id)
+            | Self::FinishCall(node_id)
+            | Self::FinishDyn(node_id)
+            | Self::ResumeBasicBlock { node_id, .. }
+            | Self::Respan { node_id, .. }
+            | Self::FinishBasicBlock(node_id) => Some(*node_id),
+            Self::EnterForest { .. } => None,
         }
     }
 }
@@ -123,12 +142,19 @@ impl<F> ContinuationStack<F> {
     pub(crate) fn new_with_source_node_id(
         program: &Program,
         source_node_id: DebugSourceNodeId,
+    ) -> Self {
+        Self::new_with_optional_source_node_id(program, Some(source_node_id))
+    }
+
+    pub(crate) fn new_with_optional_source_node_id(
+        program: &Program,
+        source_node_id: Option<DebugSourceNodeId>,
     ) -> Self {
         let mut stack = Vec::with_capacity(CONTINUATION_STACK_SIZE_HINT);
         stack.push(Continuation::StartNode(program.entrypoint()));
 
         let mut source_node_ids = Vec::with_capacity(CONTINUATION_STACK_SIZE_HINT);
-        source_node_ids.push(Some(source_node_id));
+        source_node_ids.push(source_node_id);
 
         Self {
             stack,
@@ -159,7 +185,15 @@ impl<F> ContinuationStack<F> {
     /// # Arguments
     /// * `forest` - The MAST forest to enter
     pub fn push_enter_forest(&mut self, forest: F) {
-        self.stack.push(Continuation::EnterForest(forest));
+        self.push_enter_forest_with_package_debug_info(forest, None);
+    }
+
+    pub(crate) fn push_enter_forest_with_package_debug_info(
+        &mut self,
+        forest: F,
+        package_debug_info: Option<Arc<PackageDebugInfo>>,
+    ) {
+        self.stack.push(Continuation::EnterForest { forest, package_debug_info });
         self.push_source_node_id(None);
     }
 
@@ -232,6 +266,19 @@ impl<F> ContinuationStack<F> {
         }
     }
 
+    #[cfg(any(test, feature = "testing"))]
+    pub(crate) fn start_tracking_source_nodes(
+        &mut self,
+        next_source_node_id: Option<DebugSourceNodeId>,
+    ) {
+        let mut source_node_ids = Vec::with_capacity(self.stack.len());
+        source_node_ids.resize(self.stack.len(), None);
+        if let Some(source_node_id) = source_node_ids.last_mut() {
+            *source_node_id = next_source_node_id;
+        }
+        self.source_node_ids = Some(source_node_ids);
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -260,6 +307,10 @@ impl<F> ContinuationStack<F> {
         Some((continuation, source_node_id))
     }
 
+    pub(crate) fn tracks_source_nodes(&self) -> bool {
+        self.source_node_ids.is_some()
+    }
+
     /// Returns an iterator over the continuations on the stack that will execute in the next clock
     /// cycle.
     ///
@@ -324,12 +375,15 @@ mod tests {
         // Push an incrementing continuation first (bottom of stack)
         stack.push_continuation(Continuation::StartNode(MastNodeId::new_unchecked(0)));
         // Push a non-incrementing continuation on top
-        stack.push_continuation(Continuation::EnterForest(Arc::new(MastForest::new())));
+        stack.push_continuation(Continuation::EnterForest {
+            forest: Arc::new(MastForest::new()),
+            package_debug_info: None,
+        });
 
         let result: Vec<_> = stack.iter_continuations_for_next_clock().collect();
         // Should return: EnterForest (non-incrementing), then StartNode (first incrementing)
         assert_eq!(result.len(), 2);
-        assert!(matches!(result[0], Continuation::EnterForest(_)));
+        assert!(matches!(result[0], Continuation::EnterForest { .. }));
         assert!(matches!(result[1], Continuation::StartNode(_)));
     }
 
@@ -339,14 +393,20 @@ mod tests {
         // Push an incrementing continuation first (bottom of stack)
         stack.push_continuation(Continuation::StartNode(MastNodeId::new_unchecked(0)));
         // Push two non-incrementing continuations on top
-        stack.push_continuation(Continuation::EnterForest(Arc::new(MastForest::new())));
-        stack.push_continuation(Continuation::EnterForest(Arc::new(MastForest::new())));
+        stack.push_continuation(Continuation::EnterForest {
+            forest: Arc::new(MastForest::new()),
+            package_debug_info: None,
+        });
+        stack.push_continuation(Continuation::EnterForest {
+            forest: Arc::new(MastForest::new()),
+            package_debug_info: None,
+        });
 
         let result: Vec<_> = stack.iter_continuations_for_next_clock().collect();
         // Should return: EnterForest, EnterForest, StartNode
         assert_eq!(result.len(), 3);
-        assert!(matches!(result[0], Continuation::EnterForest(_)));
-        assert!(matches!(result[1], Continuation::EnterForest(_)));
+        assert!(matches!(result[0], Continuation::EnterForest { .. }));
+        assert!(matches!(result[1], Continuation::EnterForest { .. }));
         assert!(matches!(result[2], Continuation::StartNode(_)));
     }
 }

processor/src/execution/basic_block.rs

@@ -5,6 +5,7 @@ use miden_mast_package::debug_info::DebugSourceNodeId;
 use crate::{
     BaseHost, BreakReason, Stopper,
     continuation_stack::{Continuation, ContinuationStack},
+    errors::PackageSourceDebugContext,
     execution::{
         ExecutionState, InternalBreakReason, execute_op, finalize_clock_cycle_with_continuation,
         finalize_clock_cycle_with_continuation_and_op_helpers,
@@ -281,13 +282,20 @@ where
             },
             _ => {
                 // If the operation is not an Emit, we execute it normally.
+                let source_debug_info = state.source_debug_info.clone();
+                let package_source_context = source_debug_info.as_deref().map(|debug_info| {
+                    PackageSourceDebugContext::new_optional(
+                        debug_info,
+                        state.current_source_node_id(),
+                    )
+                });
                 match execute_op(
                     state.processor,
                     op,
                     op_idx_in_block,
                     state.host,
                     state.tracer,
-                    state.package_source_context(),
+                    package_source_context,
                 ) {
                     Ok(operation_helpers) => operation_helpers,
                     Err(err) => {

processor/src/execution/dyn.rs

@@ -1,6 +1,8 @@
+use alloc::sync::Arc;
 use core::ops::ControlFlow;
 
 use miden_core::{FMP_ADDR, FMP_INIT_VALUE};
+use miden_mast_package::debug_info::{DebugSourceNodeId, PackageDebugInfo};
 
 use crate::{
     BaseHost, BreakReason, ContextId, MapExecErr, Stopper,
@@ -105,7 +107,7 @@ where
     // host).
     match current_forest.find_procedure_root(callee_hash) {
         Some(callee_id) => {
-            let source_node_id = match state.source_debug_info {
+            let source_node_id = match &state.source_debug_info {
                 Some(source_debug_info) => source_debug_info
                     .unique_source_root_for_exec_node(callee_id)
                     .unwrap_or_default(),
@@ -254,8 +256,11 @@ where
 pub fn finish_load_mast_forest_from_dyn_start<P, S, T, F>(
     root_id: MastNodeId,
     new_forest: F,
+    new_package_debug_info: Option<Arc<PackageDebugInfo>>,
+    new_source_node_id: Option<DebugSourceNodeId>,
     processor: &mut P,
     current_forest: &mut F,
+    current_package_debug_info: &mut Option<Arc<PackageDebugInfo>>,
     continuation_stack: &mut ContinuationStack<F>,
     tracer: &mut T,
     stopper: &S,
@@ -268,15 +273,19 @@ where
 {
     // Save the old forest: the continuation from start_clock_cycle references nodes in it.
     let old_forest = current_forest.clone();
+    let old_package_debug_info = current_package_debug_info.clone();
 
     // Push current forest to the continuation stack so that we can return to it
-    continuation_stack.push_enter_forest(current_forest.clone());
+    continuation_stack
+        .push_enter_forest_with_package_debug_info(current_forest.clone(), old_package_debug_info);
 
     // Push the root node of the external MAST forest onto the continuation stack.
-    continuation_stack.push_start_node(root_id);
+    continuation_stack
+        .push_with_source_node_id(Continuation::StartNode(root_id), new_source_node_id);
 
     // Set the new MAST forest as current
     *current_forest = new_forest;
+    *current_package_debug_info = new_package_debug_info;
 
     // Finalize the clock cycle corresponding to the DYN or DYNCALL operation. We pass the old
     // forest because the continuation was set during start_clock_cycle, which referenced the old

processor/src/execution/external.rs

@@ -1,6 +1,7 @@
+use alloc::sync::Arc;
 use core::ops::ControlFlow;
 
-use miden_mast_package::debug_info::DebugSourceNodeId;
+use miden_mast_package::debug_info::{DebugSourceNodeId, PackageDebugInfo};
 
 use crate::{
     BreakReason,
@@ -54,8 +55,11 @@ where
 pub fn finish_load_mast_forest_from_external<F, T>(
     resolved_node_id_new_forest: MastNodeId,
     new_mast_forest: F,
+    new_package_debug_info: Option<Arc<PackageDebugInfo>>,
+    new_source_node_id: Option<DebugSourceNodeId>,
     external_node_id_old_forest: MastNodeId,
     current_forest: &mut F,
+    current_package_debug_info: &mut Option<Arc<PackageDebugInfo>>,
     continuation_stack: &mut ContinuationStack<F>,
     tracer: &mut T,
 ) -> ControlFlow<BreakReason<F>>
@@ -83,18 +87,21 @@ where
 
     tracer.record_mast_forest_resolution(resolved_node_id_new_forest, &new_mast_forest);
 
+    let old_package_debug_info = current_package_debug_info.clone();
+
     // Push current forest to the continuation stack so that we can return to it
-    continuation_stack.push_enter_forest(old_forest.clone());
+    continuation_stack
+        .push_enter_forest_with_package_debug_info(old_forest.clone(), old_package_debug_info);
 
     // Push the root node of the external MAST forest onto the continuation stack.
-    //
-    // Caller package debug info describes the forest that contained the `External` node, not the
-    // loaded forest. The loaded root therefore starts without a source sidecar here.
-    continuation_stack
-        .push_with_source_node_id(Continuation::StartNode(resolved_node_id_new_forest), None);
+    continuation_stack.push_with_source_node_id(
+        Continuation::StartNode(resolved_node_id_new_forest),
+        new_source_node_id,
+    );
 
     // Update the current forest to the new MAST forest.
     *current_forest = new_mast_forest;
+    *current_package_debug_info = new_package_debug_info;
 
     // Note that executing an External node does not end the clock cycle, so we do not finalize the
     // clock cycle here.
@@ -137,12 +144,16 @@ mod tests {
         let mut continuation_stack =
             ContinuationStack::new_with_source_node_id(&program, caller_source_node_id);
         let mut tracer = NoopTracer;
+        let mut package_debug_info = None;
 
         let result = finish_load_mast_forest_from_external(
             target_id,
             new_mast_forest,
+            None,
+            None,
             external_id,
             &mut current_forest,
+            &mut package_debug_info,
             &mut continuation_stack,
             &mut tracer,
         );
@@ -154,7 +165,7 @@ mod tests {
         );
         assert_matches!(
             continuation_stack.pop_continuation_with_source_node_id(),
-            Some((Continuation::EnterForest(_), None))
+            Some((Continuation::EnterForest { .. }, None))
         );
         assert_matches!(
             continuation_stack.pop_continuation_with_source_node_id(),