Skip to main content

miden_testing/
mock_host.rs

1use alloc::collections::BTreeSet;
2use alloc::sync::Arc;
3use alloc::vec::Vec;
4
5use miden_processor::{
6    AdviceMutation,
7    AsyncHost,
8    BaseHost,
9    EventError,
10    FutureMaybeSend,
11    MastForest,
12    ProcessState,
13};
14use miden_protocol::transaction::TransactionEventId;
15use miden_protocol::vm::EventId;
16use miden_protocol::{CoreLibrary, Word};
17use miden_tx::TransactionExecutorHost;
18use miden_tx::auth::UnreachableAuth;
19
20use crate::TransactionContext;
21
22// MOCK HOST
23// ================================================================================================
24
25/// The [`MockHost`] wraps a [`TransactionExecutorHost`] and forwards event handling requests to it,
26/// with the difference that it only handles a subset of the events that the executor host handles.
27///
28/// Why don't we always forward requests to the executor host? In some tests, when using
29/// [`TransactionContext::execute_code`], we want to test that the transaction kernel fails
30/// with a certain error when given invalid inputs, but the event handler in the executor host would
31/// prematurely abort the transaction due to the invalid inputs. To avoid this situation, the event
32/// handler can be disabled and we can test that the transaction kernel has the expected behavior
33/// (e.g. even if the transaction host was malicious).
34///
35/// Some event handlers, such as delta or output note tracking, will similarly interfere with
36/// testing a procedure in isolation and these are also turned off in this host.
37pub(crate) struct MockHost<'store> {
38    /// The underlying [`TransactionExecutorHost`] that the mock host will forward requests to.
39    exec_host: TransactionExecutorHost<'store, 'static, TransactionContext, UnreachableAuth>,
40
41    /// The set of event IDs that the mock host will forward to the [`TransactionExecutorHost`].
42    ///
43    /// Event IDs that are not in this set are not handled. This can be useful in certain test
44    /// scenarios.
45    handled_events: BTreeSet<EventId>,
46}
47
48impl<'store> MockHost<'store> {
49    /// Returns a new [`MockHost`] instance with the provided inputs.
50    pub fn new(
51        exec_host: TransactionExecutorHost<'store, 'static, TransactionContext, UnreachableAuth>,
52    ) -> Self {
53        // CoreLibrary events are always handled.
54        let core_lib_handlers = CoreLibrary::default()
55            .handlers()
56            .into_iter()
57            .map(|(handler_event_name, _)| handler_event_name.to_event_id());
58        let mut handled_events = BTreeSet::from_iter(core_lib_handlers);
59
60        // The default set of transaction events that are always handled.
61        handled_events.extend(
62            [
63                &TransactionEventId::AccountPushProcedureIndex,
64                &TransactionEventId::LinkMapSet,
65                &TransactionEventId::LinkMapGet,
66                // TODO: It should be possible to remove this after implementing
67                // https://github.com/0xMiden/miden-base/issues/1852.
68                &TransactionEventId::EpilogueBeforeTxFeeRemovedFromAccount,
69            ]
70            .map(TransactionEventId::event_id),
71        );
72
73        Self { exec_host, handled_events }
74    }
75
76    // Adds the transaction events needed for Lazy loading to the set of handled events.
77    pub fn enable_lazy_loading(&mut self) {
78        self.handled_events.extend(
79            [
80                &TransactionEventId::AccountBeforeForeignLoad,
81                &TransactionEventId::AccountVaultBeforeGetBalance,
82                &TransactionEventId::AccountVaultBeforeHasNonFungibleAsset,
83                &TransactionEventId::AccountVaultBeforeAddAsset,
84                &TransactionEventId::AccountVaultBeforeRemoveAsset,
85                &TransactionEventId::AccountStorageBeforeSetMapItem,
86                &TransactionEventId::AccountStorageBeforeGetMapItem,
87            ]
88            .map(TransactionEventId::event_id),
89        );
90    }
91}
92
93impl<'store> BaseHost for MockHost<'store> {
94    fn get_label_and_source_file(
95        &self,
96        location: &miden_protocol::assembly::debuginfo::Location,
97    ) -> (
98        miden_protocol::assembly::debuginfo::SourceSpan,
99        Option<Arc<miden_protocol::assembly::SourceFile>>,
100    ) {
101        self.exec_host.get_label_and_source_file(location)
102    }
103}
104
105impl<'store> AsyncHost for MockHost<'store> {
106    fn get_mast_forest(&self, node_digest: &Word) -> impl FutureMaybeSend<Option<Arc<MastForest>>> {
107        self.exec_host.get_mast_forest(node_digest)
108    }
109
110    fn on_event(
111        &mut self,
112        process: &ProcessState,
113    ) -> impl FutureMaybeSend<Result<Vec<AdviceMutation>, EventError>> {
114        let event_id = EventId::from_felt(process.get_stack_item(0));
115
116        async move {
117            // If the host should handle the event, delegate to the tx executor host.
118            if self.handled_events.contains(&event_id) {
119                self.exec_host.on_event(process).await
120            } else {
121                Ok(Vec::new())
122            }
123        }
124    }
125}