Skip to main content

miden_debug/debug/
stacktrace.rs

1use std::{
2    borrow::Cow,
3    cell::{OnceCell, RefCell},
4    collections::{BTreeMap, BTreeSet, VecDeque},
5    fmt,
6    path::Path,
7    rc::Rc,
8    sync::Arc,
9};
10
11use miden_core::operations::AssemblyOp;
12use miden_debug_types::{Location, SourceFile, SourceManager, SourceManagerExt, SourceSpan};
13use miden_processor::{ContextId, operation::Operation, trace::RowIndex};
14
15use crate::exec::TraceEvent;
16
17pub struct StepInfo<'a> {
18    pub op: Option<Operation>,
19    pub asmop: Option<&'a AssemblyOp>,
20    pub clk: RowIndex,
21    pub ctx: ContextId,
22}
23
24#[derive(Debug, Clone)]
25struct SpanContext {
26    frame_index: usize,
27    location: Option<Location>,
28}
29
30pub struct CallStack {
31    trace_events: Rc<RefCell<BTreeMap<RowIndex, TraceEvent>>>,
32    contexts: BTreeSet<Rc<str>>,
33    frames: Vec<CallFrame>,
34    block_stack: Vec<Option<SpanContext>>,
35}
36impl CallStack {
37    pub fn new(trace_events: Rc<RefCell<BTreeMap<RowIndex, TraceEvent>>>) -> Self {
38        Self {
39            trace_events,
40            contexts: BTreeSet::default(),
41            frames: vec![],
42            block_stack: vec![],
43        }
44    }
45
46    pub fn stacktrace<'a>(
47        &'a self,
48        recent: &'a VecDeque<Operation>,
49        source_manager: &'a dyn SourceManager,
50    ) -> StackTrace<'a> {
51        StackTrace::new(self, recent, source_manager)
52    }
53
54    pub fn current_frame(&self) -> Option<&CallFrame> {
55        self.frames.last()
56    }
57
58    pub fn current_frame_mut(&mut self) -> Option<&mut CallFrame> {
59        self.frames.last_mut()
60    }
61
62    pub fn frames(&self) -> &[CallFrame] {
63        self.frames.as_slice()
64    }
65
66    /// Updates the call stack from `info`
67    ///
68    /// Returns the call frame exited this cycle, if any
69    pub fn next(&mut self, info: &StepInfo<'_>) -> Option<CallFrame> {
70        if let Some(op) = info.op {
71            // Get the current procedure name context, if available
72            let procedure = info.asmop.map(|op| self.cache_procedure_name(op.context_name()));
73
74            // Handle trace events for this cycle
75            let event = self.trace_events.borrow().get(&info.clk).copied();
76            log::trace!("handling {op} at cycle {}: {:?}", info.clk, &event);
77            let popped_frame = self.handle_trace_event(event, procedure.as_ref());
78            let is_frame_end = popped_frame.is_some();
79
80            // These ops we do not record in call frame details
81            let ignore = matches!(
82                op,
83                Operation::Join
84                    | Operation::Split
85                    | Operation::Span
86                    | Operation::Respan
87                    | Operation::End
88            );
89
90            // Manage block stack
91            match op {
92                Operation::Span => {
93                    if let Some(asmop) = info.asmop {
94                        log::debug!("{asmop:#?}");
95                        self.block_stack.push(Some(SpanContext {
96                            frame_index: self.frames.len().saturating_sub(1),
97                            location: asmop.location().cloned(),
98                        }));
99                    } else {
100                        self.block_stack.push(None);
101                    }
102                }
103                Operation::End => {
104                    self.block_stack.pop();
105                }
106                Operation::Join | Operation::Split => {
107                    self.block_stack.push(None);
108                }
109                _ => (),
110            }
111
112            if ignore || is_frame_end {
113                return popped_frame;
114            }
115
116            // Attempt to supply procedure context from the current span context, if needed +
117            // available
118            let (procedure, asmop) = match procedure {
119                proc @ Some(_) => (proc, info.asmop.map(Cow::Borrowed)),
120                None => match self.block_stack.last() {
121                    Some(Some(span_ctx)) => {
122                        let proc =
123                            self.frames.get(span_ctx.frame_index).and_then(|f| f.procedure.clone());
124                        let asmop_cow = info.asmop.map(Cow::Borrowed).or_else(|| {
125                            let context_name = proc.as_deref().unwrap_or("<unknown>").to_string();
126                            let raw_asmop = AssemblyOp::new(
127                                span_ctx.location.clone(),
128                                context_name,
129                                1,
130                                op.to_string(),
131                            );
132                            Some(Cow::Owned(raw_asmop))
133                        });
134                        (proc, asmop_cow)
135                    }
136                    _ => (None, info.asmop.map(Cow::Borrowed)),
137                },
138            };
139
140            // Use the current frame's procedure context, if no other more precise context is
141            // available
142            let procedure =
143                procedure.or_else(|| self.frames.last().and_then(|f| f.procedure.clone()));
144
145            // Do we have a frame? If not, create one
146            if self.frames.is_empty() {
147                self.frames.push(CallFrame::new(procedure.clone()));
148            }
149
150            let current_frame = self.frames.last_mut().unwrap();
151
152            // Does the current frame have a procedure context/location? Use the one from this op if
153            // so
154            let procedure_context_updated =
155                current_frame.procedure.is_none() && procedure.is_some();
156            if procedure_context_updated {
157                current_frame.procedure.clone_from(&procedure);
158            }
159
160            // Push op into call frame if this is any op other than `nop` or frame setup
161            if !matches!(op, Operation::Noop) {
162                let cycle_idx = info.asmop.map(|a| a.num_cycles()).unwrap_or(1);
163                current_frame.push(op, cycle_idx, asmop.as_deref());
164            }
165
166            // Check if we should also update the caller frame's exec detail
167            let num_frames = self.frames.len();
168            if procedure_context_updated && num_frames > 1 {
169                let caller_frame = &mut self.frames[num_frames - 2];
170                if let Some(OpDetail::Exec { callee }) = caller_frame.context.back_mut()
171                    && callee.is_none()
172                {
173                    *callee = procedure;
174                }
175            }
176        }
177
178        None
179    }
180
181    // Get or cache procedure name/context as `Rc<str>`
182    fn cache_procedure_name(&mut self, context_name: &str) -> Rc<str> {
183        match self.contexts.get(context_name) {
184            Some(name) => Rc::clone(name),
185            None => {
186                let name = Rc::from(context_name.to_string().into_boxed_str());
187                self.contexts.insert(Rc::clone(&name));
188                name
189            }
190        }
191    }
192
193    fn handle_trace_event(
194        &mut self,
195        event: Option<TraceEvent>,
196        procedure: Option<&Rc<str>>,
197    ) -> Option<CallFrame> {
198        // Do we need to handle any frame events?
199        if let Some(event) = event {
200            match event {
201                TraceEvent::FrameStart => {
202                    // Record the fact that we exec'd a new procedure in the op context
203                    if let Some(current_frame) = self.frames.last_mut() {
204                        current_frame.push_exec(procedure.cloned());
205                    }
206                    // Push a new frame
207                    self.frames.push(CallFrame::new(procedure.cloned()));
208                }
209                TraceEvent::Unknown(code) => log::debug!("unknown trace event: {code}"),
210                TraceEvent::FrameEnd => {
211                    return self.frames.pop();
212                }
213                _ => (),
214            }
215        }
216        None
217    }
218}
219
220pub struct CallFrame {
221    procedure: Option<Rc<str>>,
222    context: VecDeque<OpDetail>,
223    display_name: std::cell::OnceCell<Rc<str>>,
224    finishing: bool,
225}
226impl CallFrame {
227    pub fn new(procedure: Option<Rc<str>>) -> Self {
228        Self {
229            procedure,
230            context: Default::default(),
231            display_name: Default::default(),
232            finishing: false,
233        }
234    }
235
236    pub fn procedure(&self, strip_prefix: &str) -> Option<Rc<str>> {
237        self.procedure.as_ref()?;
238        let name = self.display_name.get_or_init(|| {
239            let name = self.procedure.as_deref().unwrap();
240            let name = match name.split_once("::") {
241                Some((module, rest)) if module == strip_prefix => demangle(rest),
242                _ => demangle(name),
243            };
244            Rc::from(name.into_boxed_str())
245        });
246        Some(Rc::clone(name))
247    }
248
249    pub fn push_exec(&mut self, callee: Option<Rc<str>>) {
250        if self.context.len() == 5 {
251            self.context.pop_front();
252        }
253
254        self.context.push_back(OpDetail::Exec { callee });
255    }
256
257    pub fn push(&mut self, opcode: Operation, cycle_idx: u8, op: Option<&AssemblyOp>) {
258        if cycle_idx > 1 {
259            // Should we ignore this op?
260            let skip = self.context.back().map(|detail| matches!(detail, OpDetail::Full { op, .. } | OpDetail::Basic { op } if op == &opcode)).unwrap_or(false);
261            if skip {
262                return;
263            }
264        }
265
266        if self.context.len() == 5 {
267            self.context.pop_front();
268        }
269
270        match op {
271            Some(op) => {
272                let location = op.location().cloned();
273                self.context.push_back(OpDetail::Full {
274                    op: opcode,
275                    location,
276                    resolved: Default::default(),
277                });
278            }
279            None => {
280                // If this instruction does not have a location, inherit the location
281                // of the previous op in the frame, if one is present
282                if let Some(loc) = self.context.back().map(|op| op.location().cloned()) {
283                    self.context.push_back(OpDetail::Full {
284                        op: opcode,
285                        location: loc,
286                        resolved: Default::default(),
287                    });
288                } else {
289                    self.context.push_back(OpDetail::Basic { op: opcode });
290                }
291            }
292        }
293    }
294
295    pub fn last_location(&self) -> Option<&Location> {
296        match self.context.back() {
297            Some(OpDetail::Full { location, .. }) => {
298                let loc = location.as_ref();
299                if loc.is_none() {
300                    dbg!(&self.context);
301                }
302                loc
303            }
304            Some(OpDetail::Basic { .. }) => None,
305            Some(OpDetail::Exec { .. }) => {
306                let op = self.context.iter().rev().nth(1)?;
307                op.location()
308            }
309            None => None,
310        }
311    }
312
313    pub fn last_resolved(&self, source_manager: &dyn SourceManager) -> Option<&ResolvedLocation> {
314        // Search through context in reverse order to find the most recent op with a resolvable
315        // location.
316        for op in self.context.iter().rev() {
317            if let Some(resolved) = op.resolve(source_manager) {
318                return Some(resolved);
319            }
320        }
321        None
322    }
323
324    pub fn recent(&self) -> &VecDeque<OpDetail> {
325        &self.context
326    }
327
328    #[inline(always)]
329    pub fn should_break_on_exit(&self) -> bool {
330        self.finishing
331    }
332
333    #[inline(always)]
334    pub fn break_on_exit(&mut self) {
335        self.finishing = true;
336    }
337}
338
339#[derive(Debug, Clone)]
340pub enum OpDetail {
341    Full {
342        op: Operation,
343        location: Option<Location>,
344        resolved: OnceCell<Option<ResolvedLocation>>,
345    },
346    Exec {
347        callee: Option<Rc<str>>,
348    },
349    Basic {
350        op: Operation,
351    },
352}
353impl OpDetail {
354    pub fn callee(&self, strip_prefix: &str) -> Option<Box<str>> {
355        match self {
356            Self::Exec { callee: None } => Some(Box::from("<unknown>")),
357            Self::Exec {
358                callee: Some(callee),
359            } => {
360                let name = match callee.split_once("::") {
361                    Some((module, rest)) if module == strip_prefix => demangle(rest),
362                    _ => demangle(callee),
363                };
364                Some(name.into_boxed_str())
365            }
366            _ => None,
367        }
368    }
369
370    pub fn display(&self) -> String {
371        match self {
372            Self::Full { op, .. } | Self::Basic { op } => format!("{op}"),
373            Self::Exec {
374                callee: Some(callee),
375            } => format!("exec.{callee}"),
376            Self::Exec { callee: None } => "exec.<unavailable>".to_string(),
377        }
378    }
379
380    pub fn opcode(&self) -> Operation {
381        match self {
382            Self::Full { op, .. } | Self::Basic { op } => *op,
383            Self::Exec { .. } => panic!("no opcode associated with execs"),
384        }
385    }
386
387    pub fn location(&self) -> Option<&Location> {
388        match self {
389            Self::Full { location, .. } => location.as_ref(),
390            Self::Basic { .. } | Self::Exec { .. } => None,
391        }
392    }
393
394    pub fn resolve(&self, source_manager: &dyn SourceManager) -> Option<&ResolvedLocation> {
395        match self {
396            Self::Full {
397                location: Some(loc),
398                resolved,
399                ..
400            } => resolved
401                .get_or_init(|| {
402                    let path = Path::new(loc.uri().as_str());
403                    let source_file = if path.exists() {
404                        source_manager.load_file(path).ok()?
405                    } else {
406                        source_manager.get_by_uri(loc.uri())?
407                    };
408                    let span = SourceSpan::new(source_file.id(), loc.start..loc.end);
409                    let file_line_col = source_file.location(span);
410                    Some(ResolvedLocation {
411                        source_file,
412                        line: file_line_col.line.to_u32(),
413                        col: file_line_col.column.to_u32(),
414                        span,
415                    })
416                })
417                .as_ref(),
418            _ => None,
419        }
420    }
421}
422
423#[derive(Debug, Clone)]
424pub struct ResolvedLocation {
425    pub source_file: Arc<SourceFile>,
426    // TODO(fabrio): Use LineNumber and ColumnNumber instead of raw `u32`.
427    pub line: u32,
428    pub col: u32,
429    pub span: SourceSpan,
430}
431impl fmt::Display for ResolvedLocation {
432    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
433        write!(f, "{}:{}:{}", self.source_file.uri().as_str(), self.line, self.col)
434    }
435}
436
437pub struct CurrentFrame {
438    pub procedure: Option<Rc<str>>,
439    pub location: Option<ResolvedLocation>,
440}
441
442pub struct StackTrace<'a> {
443    callstack: &'a CallStack,
444    recent: &'a VecDeque<Operation>,
445    source_manager: &'a dyn SourceManager,
446    current_frame: Option<CurrentFrame>,
447}
448
449impl<'a> StackTrace<'a> {
450    pub fn new(
451        callstack: &'a CallStack,
452        recent: &'a VecDeque<Operation>,
453        source_manager: &'a dyn SourceManager,
454    ) -> Self {
455        let current_frame = callstack.current_frame().map(|frame| {
456            let location = frame.last_resolved(source_manager).cloned();
457            let procedure = frame.procedure("");
458            CurrentFrame {
459                procedure,
460                location,
461            }
462        });
463        Self {
464            callstack,
465            recent,
466            source_manager,
467            current_frame,
468        }
469    }
470
471    pub fn current_frame(&self) -> Option<&CurrentFrame> {
472        self.current_frame.as_ref()
473    }
474}
475
476impl fmt::Display for StackTrace<'_> {
477    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
478        use std::fmt::Write;
479
480        let num_frames = self.callstack.frames.len();
481
482        writeln!(f, "\nStack Trace:")?;
483
484        for (i, frame) in self.callstack.frames.iter().enumerate() {
485            let is_top = i + 1 == num_frames;
486            let name = frame.procedure("");
487            let name = name.as_deref().unwrap_or("<unknown>");
488            if is_top {
489                write!(f, " `-> {name}")?;
490            } else {
491                write!(f, " |-> {name}")?;
492            }
493            if let Some(resolved) = frame.last_resolved(self.source_manager) {
494                write!(f, " in {resolved}")?;
495            } else {
496                write!(f, " in <unavailable>")?;
497            }
498            if is_top {
499                // Print op context
500                let context_size = frame.context.len();
501                writeln!(f, ":\n\nLast {context_size} Instructions (of current frame):")?;
502                for (i, op) in frame.context.iter().enumerate() {
503                    let is_last = i + 1 == context_size;
504                    if let Some(callee) = op.callee("") {
505                        write!(f, " |   exec.{callee}")?;
506                    } else {
507                        write!(f, " |   {}", &op.opcode())?;
508                    }
509                    if is_last {
510                        writeln!(f, "\n `-> <error occured here>")?;
511                    } else {
512                        f.write_char('\n')?;
513                    }
514                }
515
516                let context_size = self.recent.len();
517                writeln!(f, "\n\nLast {context_size} Instructions (any frame):")?;
518                for (i, op) in self.recent.iter().enumerate() {
519                    let is_last = i + 1 == context_size;
520                    if is_last {
521                        writeln!(f, " |   {}", &op)?;
522                        writeln!(f, " `-> <error occured here>")?;
523                    } else {
524                        writeln!(f, " |   {}", &op)?;
525                    }
526                }
527            } else {
528                f.write_char('\n')?;
529            }
530        }
531
532        Ok(())
533    }
534}
535
536fn demangle(name: &str) -> String {
537    let mut input = name.as_bytes();
538    let mut demangled = Vec::with_capacity(input.len() * 2);
539    rustc_demangle::demangle_stream(&mut input, &mut demangled, /* include_hash= */ false)
540        .expect("failed to write demangled identifier");
541    String::from_utf8(demangled).expect("demangled identifier contains invalid utf-8")
542}