midenc_debug/exec/
host.rs

1use std::{collections::BTreeMap, num::NonZeroU32, sync::Arc};
2
3use miden_core::crypto::hash::RpoDigest;
4use miden_processor::{
5    AdviceInputs, AdviceProvider, ExecutionError, Host, KvMap, MastForest, MastForestStore,
6    MemAdviceProvider, 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        // Extract and load the advice map from the forest before putting it into the store.
55        let advice_map = forest.advice_map();
56        for (digest, values) in advice_map.iter() {
57            let key = digest.into();
58            self.adv_provider.insert_into_map(key, values.clone());
59        }
60        self.store.insert(forest);
61    }
62}
63
64impl Host for DebuggerHost {
65    type AdviceProvider = MemAdviceProvider;
66
67    fn get_mast_forest(&self, node_digest: &RpoDigest) -> Option<Arc<MastForest>> {
68        self.store.get(node_digest)
69    }
70
71    fn on_trace(&mut self, process: ProcessState, trace_id: u32) -> Result<(), ExecutionError> {
72        let event = TraceEvent::from(trace_id);
73        let clk = process.clk();
74        if let Some(handlers) = self.tracing_callbacks.get_mut(&trace_id) {
75            for handler in handlers.iter_mut() {
76                handler(clk, event);
77            }
78        }
79        Ok(())
80    }
81
82    fn on_assert_failed(&mut self, process: ProcessState, err_code: miden_core::Felt) {
83        let clk = process.clk();
84        if let Some(handler) = self.on_assert_failed.as_mut() {
85            // TODO: We're truncating the error code here, but we may need to handle the full range
86            handler(clk, TraceEvent::AssertionFailed(NonZeroU32::new(err_code.as_int() as u32)));
87        }
88    }
89
90    fn advice_provider(&self) -> &Self::AdviceProvider {
91        &self.adv_provider
92    }
93
94    fn advice_provider_mut(&mut self) -> &mut Self::AdviceProvider {
95        &mut self.adv_provider
96    }
97}