Skip to main content

miden_debug_engine/exec/
host.rs

1use std::{
2    collections::{BTreeMap, VecDeque},
3    num::NonZeroU32,
4    sync::Arc,
5};
6
7use miden_assembly::SourceManager;
8use miden_core::{
9    Word,
10    events::{EventId, EventName},
11};
12use miden_debug_types::{Location, SourceFile, SourceSpan};
13use miden_processor::{
14    ExecutionError, FutureMaybeSend, Host, MastForestStore, MemMastForestStore, ProcessorState,
15    TraceError,
16    advice::AdviceMutation,
17    event::{EventError, EventHandler, EventHandlerRegistry},
18    mast::MastForest,
19};
20
21use super::{TraceEvent, TraceHandler};
22
23/// This is an implementation of [Host] which is essentially [miden_processor::DefaultHost],
24/// but extended with additional functionality for debugging, in particular it manages trace
25/// events that record the entry or exit of a procedure call frame.
26pub struct DebuggerHost<S: SourceManager + ?Sized> {
27    store: MemMastForestStore,
28    event_handlers: EventHandlerRegistry,
29    tracing_callbacks: BTreeMap<u32, Vec<Box<TraceHandler>>>,
30    on_assert_failed: Option<Box<TraceHandler>>,
31    source_manager: Arc<S>,
32    event_replay: VecDeque<Vec<AdviceMutation>>,
33}
34impl<S> DebuggerHost<S>
35where
36    S: SourceManager + ?Sized,
37{
38    /// Construct a new instance of [DebuggerHost] with the given source manager.
39    pub fn new(source_manager: Arc<S>) -> Self {
40        Self {
41            store: Default::default(),
42            event_handlers: EventHandlerRegistry::default(),
43            tracing_callbacks: Default::default(),
44            on_assert_failed: None,
45            source_manager,
46            event_replay: VecDeque::new(),
47        }
48    }
49
50    /// Set the event replay queue.
51    ///
52    /// When non-empty, `on_event()` will pop mutations from this queue instead of
53    /// returning empty results. This is used for transaction debugging where events
54    /// were recorded during a prior execution.
55    pub fn set_event_replay(&mut self, events: VecDeque<Vec<AdviceMutation>>) {
56        self.event_replay = events;
57    }
58
59    /// Register a trace handler for `event`
60    pub fn register_trace_handler<F>(&mut self, event: TraceEvent, callback: F)
61    where
62        F: FnMut(&ProcessorState<'_>, TraceEvent) + 'static,
63    {
64        let key = match event {
65            TraceEvent::AssertionFailed(None) => u32::MAX,
66            ev => ev.into(),
67        };
68        self.tracing_callbacks.entry(key).or_default().push(Box::new(callback));
69    }
70
71    /// Register a handler to be called when an assertion in the VM fails
72    pub fn register_assert_failed_tracer<F>(&mut self, callback: F)
73    where
74        F: FnMut(&ProcessorState<'_>, TraceEvent) + 'static,
75    {
76        self.on_assert_failed = Some(Box::new(callback));
77    }
78
79    /// Invoke the assert-failed handler, if registered.
80    ///
81    /// This is called externally when `step()` returns an assertion error, since
82    /// `on_assert_failed` no longer exists on the Host trait in 0.21.
83    pub fn handle_assert_failed(
84        &mut self,
85        process: &ProcessorState<'_>,
86        err_code: Option<NonZeroU32>,
87    ) {
88        if let Some(handler) = self.on_assert_failed.as_mut() {
89            handler(process, TraceEvent::AssertionFailed(err_code));
90        }
91    }
92
93    /// Load `forest` into the MAST store for this host
94    pub fn load_mast_forest(&mut self, forest: Arc<MastForest>) {
95        self.store.insert(forest);
96    }
97
98    /// Registers an event handler for use during program execution.
99    pub fn register_event_handler(
100        &mut self,
101        event: EventName,
102        handler: Arc<dyn EventHandler>,
103    ) -> Result<(), ExecutionError> {
104        self.event_handlers.register(event, handler)
105    }
106}
107
108impl<S> Host for DebuggerHost<S>
109where
110    S: SourceManager + ?Sized,
111{
112    fn get_label_and_source_file(
113        &self,
114        location: &Location,
115    ) -> (SourceSpan, Option<Arc<SourceFile>>) {
116        let maybe_file = self.source_manager.get_by_uri(location.uri());
117        let span = self.source_manager.location_to_span(location.clone()).unwrap_or_default();
118        (span, maybe_file)
119    }
120
121    fn get_mast_forest(&self, node_digest: &Word) -> impl FutureMaybeSend<Option<Arc<MastForest>>> {
122        std::future::ready(self.store.get(node_digest))
123    }
124
125    fn on_event(
126        &mut self,
127        process: &ProcessorState<'_>,
128    ) -> impl FutureMaybeSend<Result<Vec<AdviceMutation>, EventError>> {
129        if !self.event_replay.is_empty() {
130            let mutations = self.event_replay.pop_front().unwrap_or_default();
131            return std::future::ready(Ok(mutations));
132        }
133
134        let event_id = EventId::from_felt(process.get_stack_item(0));
135        let result = match self.event_handlers.handle_event(event_id, process) {
136            Ok(Some(mutations)) => Ok(mutations),
137            Ok(None) => {
138                #[derive(Debug, thiserror::Error)]
139                #[error("no event handler registered")]
140                struct UnhandledEvent;
141
142                Err(UnhandledEvent.into())
143            }
144            Err(err) => Err(err),
145        };
146        std::future::ready(result)
147    }
148
149    fn on_trace(&mut self, process: &ProcessorState<'_>, trace_id: u32) -> Result<(), TraceError> {
150        let event = TraceEvent::from(trace_id);
151        if let Some(handlers) = self.tracing_callbacks.get_mut(&trace_id) {
152            for handler in handlers.iter_mut() {
153                handler(process, event);
154            }
155        }
156        Ok(())
157    }
158
159    fn resolve_event(&self, event_id: EventId) -> Option<&EventName> {
160        self.event_handlers.resolve_event(event_id)
161    }
162}