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 pub fn next(&mut self, info: &StepInfo<'_>) -> Option<CallFrame> {
70 if let Some(op) = info.op {
71 let procedure = info.asmop.map(|op| self.cache_procedure_name(op.context_name()));
73
74 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 let ignore = matches!(
82 op,
83 Operation::Join
84 | Operation::Split
85 | Operation::Span
86 | Operation::Respan
87 | Operation::End
88 );
89
90 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 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 let procedure =
143 procedure.or_else(|| self.frames.last().and_then(|f| f.procedure.clone()));
144
145 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 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 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 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 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 if let Some(event) = event {
200 match event {
201 TraceEvent::FrameStart => {
202 if let Some(current_frame) = self.frames.last_mut() {
204 current_frame.push_exec(procedure.cloned());
205 }
206 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 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 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 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 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 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, false)
540 .expect("failed to write demangled identifier");
541 String::from_utf8(demangled).expect("demangled identifier contains invalid utf-8")
542}