miden_tx/host/
script_mast_forest_store.rs

1use alloc::collections::BTreeMap;
2use alloc::sync::Arc;
3
4use miden_processor::MastForestStore;
5use miden_protocol::Word;
6use miden_protocol::assembly::mast::MastForest;
7use miden_protocol::note::NoteScript;
8use miden_protocol::transaction::TransactionScript;
9use miden_protocol::vm::AdviceMap;
10
11/// Stores the MAST forests for a set of scripts (both note scripts and transaction scripts).
12///
13/// A [ScriptMastForestStore] is meant to exclusively store MAST forests related to both
14/// transaction and input note scripts.
15#[derive(Debug, Clone, Default)]
16pub struct ScriptMastForestStore {
17    mast_forests: BTreeMap<Word, Arc<MastForest>>,
18    advice_map: AdviceMap,
19}
20
21impl ScriptMastForestStore {
22    /// Creates a new [ScriptMastForestStore].
23    pub fn new(
24        tx_script: Option<&TransactionScript>,
25        note_scripts: impl Iterator<Item = impl AsRef<NoteScript>>,
26    ) -> Self {
27        let mut mast_store = ScriptMastForestStore {
28            mast_forests: BTreeMap::new(),
29            advice_map: AdviceMap::default(),
30        };
31
32        for note_script in note_scripts {
33            mast_store.insert(note_script.as_ref().mast());
34        }
35
36        if let Some(tx_script) = tx_script {
37            mast_store.insert(tx_script.mast());
38        }
39        mast_store
40    }
41
42    /// Registers all procedures of the provided [MastForest] with this store.
43    fn insert(&mut self, mast_forest: Arc<MastForest>) {
44        // only register procedures that are local to this forest
45        for proc_digest in mast_forest.local_procedure_digests() {
46            self.mast_forests.insert(proc_digest, mast_forest.clone());
47        }
48
49        // collect advice data from the forest
50        for (key, values) in mast_forest.advice_map().clone() {
51            self.advice_map.insert((*key).into(), values);
52        }
53    }
54
55    /// Returns a reference to the advice data collected from all forests.
56    pub fn advice_map(&self) -> &AdviceMap {
57        &self.advice_map
58    }
59}
60
61// MAST FOREST STORE IMPLEMENTATION
62// ================================================================================================
63
64impl MastForestStore for ScriptMastForestStore {
65    fn get(&self, procedure_root: &Word) -> Option<Arc<MastForest>> {
66        self.mast_forests.get(procedure_root).cloned()
67    }
68}