use alloc::collections::BTreeMap;
use alloc::sync::Arc;
use miden_processor::MastForestStore;
use miden_protocol::account::AccountCode;
use miden_protocol::assembly::mast::MastForest;
use miden_protocol::transaction::TransactionKernel;
use miden_protocol::utils::sync::RwLock;
use miden_protocol::{CoreLibrary, ProtocolLib, Word};
use miden_standards::StandardsLib;
pub struct TransactionMastStore {
mast_forests: RwLock<BTreeMap<Word, Arc<MastForest>>>,
}
#[allow(clippy::new_without_default)]
impl TransactionMastStore {
pub fn new() -> Self {
let mast_forests = RwLock::new(BTreeMap::new());
let store = Self { mast_forests };
let kernels_forest = TransactionKernel::kernel().mast_forest().clone();
store.insert(kernels_forest);
let miden_core_lib_forest = CoreLibrary::default().mast_forest().clone();
store.insert(miden_core_lib_forest);
let protocol_lib_forest = ProtocolLib::default().mast_forest().clone();
store.insert(protocol_lib_forest);
let standards_lib_forest = StandardsLib::default().mast_forest().clone();
store.insert(standards_lib_forest);
store
}
pub fn insert(&self, mast_forest: Arc<MastForest>) {
let mut mast_forests = self.mast_forests.write();
for proc_digest in mast_forest.local_procedure_digests() {
mast_forests.insert(proc_digest, mast_forest.clone());
}
}
pub fn load_account_code(&self, code: &AccountCode) {
self.insert(code.mast().clone());
}
}
impl MastForestStore for TransactionMastStore {
fn get(&self, procedure_root: &Word) -> Option<Arc<MastForest>> {
self.mast_forests.read().get(procedure_root).cloned()
}
}