midenc_debug/exec/
host.rs

1use std::{collections::BTreeMap, num::NonZeroU32, sync::Arc};
2
3use miden_core::crypto::hash::RpoDigest;
4use miden_processor::{
5    AdviceProvider, ExecutionError, Host, MastForest, MastForestStore, MemAdviceProvider,
6    MemMastForestStore, ProcessState, RowIndex,
7};
8
9use super::{TraceEvent, TraceHandler};
10
11/// This is an implementation of [Host] which is essentially [miden_processor::DefaultHost],
12/// but extended with additional functionality for debugging, in particular it manages trace
13/// events that record the entry or exit of a procedure call frame.
14#[derive(Default)]
15pub struct DebuggerHost {
16    adv_provider: MemAdviceProvider,
17    store: MemMastForestStore,
18    tracing_callbacks: BTreeMap<u32, Vec<Box<TraceHandler>>>,
19    on_assert_failed: Option<Box<TraceHandler>>,
20}
21impl DebuggerHost {
22    /// Construct a new instance of [DebuggerHost] with the given advice provider.
23    pub fn new(adv_provider: MemAdviceProvider) -> Self {
24        Self {
25            adv_provider,
26            store: Default::default(),
27            tracing_callbacks: Default::default(),
28            on_assert_failed: None,
29        }
30    }
31
32    /// Register a trace handler for `event`
33    pub fn register_trace_handler<F>(&mut self, event: TraceEvent, callback: F)
34    where
35        F: FnMut(RowIndex, TraceEvent) + 'static,
36    {
37        let key = match event {
38            TraceEvent::AssertionFailed(None) => u32::MAX,
39            ev => ev.into(),
40        };
41        self.tracing_callbacks.entry(key).or_default().push(Box::new(callback));
42    }
43
44    /// Register a handler to be called when an assertion in the VM fails
45    pub fn register_assert_failed_tracer<F>(&mut self, callback: F)
46    where
47        F: FnMut(RowIndex, TraceEvent) + 'static,
48    {
49        self.on_assert_failed = Some(Box::new(callback));
50    }
51
52    /// Load `forest` into the MAST store for this host
53    pub fn load_mast_forest(&mut self, forest: Arc<MastForest>) {
54        self.store.insert(forest);
55    }
56}
57
58impl Host for DebuggerHost {
59    type AdviceProvider = MemAdviceProvider;
60
61    fn get_mast_forest(&self, node_digest: &RpoDigest) -> Option<Arc<MastForest>> {
62        self.store.get(node_digest)
63    }
64
65    fn on_trace(&mut self, process: ProcessState, trace_id: u32) -> Result<(), ExecutionError> {
66        let event = TraceEvent::from(trace_id);
67        let clk = process.clk();
68        if let Some(handlers) = self.tracing_callbacks.get_mut(&trace_id) {
69            for handler in handlers.iter_mut() {
70                handler(clk, event);
71            }
72        }
73        Ok(())
74    }
75
76    fn on_assert_failed(&mut self, process: ProcessState, err_code: miden_core::Felt) {
77        let clk = process.clk();
78        if let Some(handler) = self.on_assert_failed.as_mut() {
79            // TODO: We're truncating the error code here, but we may need to handle the full range
80            handler(clk, TraceEvent::AssertionFailed(NonZeroU32::new(err_code.as_int() as u32)));
81        }
82    }
83
84    fn advice_provider(&self) -> &Self::AdviceProvider {
85        &self.adv_provider
86    }
87
88    fn advice_provider_mut(&mut self) -> &mut Self::AdviceProvider {
89        &mut self.adv_provider
90    }
91}