miden_client/store/
smt_forest.rs

1use alloc::collections::BTreeMap;
2use alloc::vec::Vec;
3
4use miden_protocol::account::{
5    AccountId,
6    AccountStorage,
7    StorageMap,
8    StorageMapKey,
9    StorageMapKeyHash,
10    StorageMapWitness,
11    StorageSlotContent,
12};
13use miden_protocol::asset::{Asset, AssetId, AssetVault, AssetWitness};
14use miden_protocol::crypto::merkle::EmptySubtreeRoots;
15use miden_protocol::crypto::merkle::smt::{SMT_DEPTH, Smt, SmtForest};
16use miden_protocol::{EMPTY_WORD, Word};
17
18use super::StoreError;
19
20/// Thin wrapper around `SmtForest` for account vault/storage proofs and updates.
21///
22/// Tracks current SMT roots per account with reference counting to safely pop
23/// roots from the underlying forest when no account references them anymore.
24/// Supports staged updates for transaction rollback via a pending roots stack.
25#[derive(Debug, Default, Clone, Eq, PartialEq)]
26pub struct AccountSmtForest {
27    forest: SmtForest,
28    /// Current roots per account (vault root + storage map roots).
29    account_roots: BTreeMap<AccountId, Vec<Word>>,
30    /// Stack of old roots saved during staging, awaiting commit or undo.
31    pending_old_roots: BTreeMap<AccountId, Vec<Vec<Word>>>,
32    /// Reference count for each SMT root across all accounts.
33    root_refcounts: BTreeMap<Word, usize>,
34}
35
36impl AccountSmtForest {
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    // READERS
42    // --------------------------------------------------------------------------------------------
43
44    /// Returns the current roots for an account.
45    pub fn get_roots(&self, account_id: &AccountId) -> Option<&Vec<Word>> {
46        self.account_roots.get(account_id)
47    }
48
49    /// Retrieves the vault asset and its witness for a specific vault key.
50    pub fn get_asset_and_witness(
51        &self,
52        vault_root: Word,
53        vault_id: AssetId,
54    ) -> Result<(Asset, AssetWitness), StoreError> {
55        let hashed_key: Word = vault_id.hash().into();
56        let proof = self.forest.open(vault_root, hashed_key)?;
57        let asset_word = proof
58            .get(&hashed_key)
59            .ok_or(StoreError::VaultKeyNotTracked(vault_id, hashed_key))?;
60        if asset_word == EMPTY_WORD {
61            return Err(StoreError::VaultKeyNotTracked(vault_id, hashed_key));
62        }
63
64        let asset = Asset::from_id_and_value(vault_id, asset_word)?;
65        let witness = AssetWitness::new(proof, [vault_id])?;
66        Ok((asset, witness))
67    }
68
69    /// Retrieves the storage map witness for a specific map item.
70    pub fn get_storage_map_item_witness(
71        &self,
72        map_root: Word,
73        key: StorageMapKey,
74    ) -> Result<StorageMapWitness, StoreError> {
75        let hashed_key = key.hash();
76        let proof = self.forest.open(map_root, Word::from(hashed_key)).map_err(StoreError::from)?;
77        Ok(StorageMapWitness::new(proof, [key])?)
78    }
79
80    // ROOT LIFECYCLE
81    // --------------------------------------------------------------------------------------------
82
83    /// Stages new roots for an account, saving old roots for potential rollback.
84    ///
85    /// The old roots are pushed onto a pending stack and their refcounts are preserved.
86    /// Call [`Self::commit_roots`] to release old roots or [`Self::discard_roots`] to
87    /// restore them.
88    pub fn stage_roots(&mut self, account_id: AccountId, new_roots: Vec<Word>) {
89        increment_refcounts(&mut self.root_refcounts, &new_roots);
90        if let Some(old_roots) = self.account_roots.insert(account_id, new_roots) {
91            self.pending_old_roots.entry(account_id).or_default().push(old_roots);
92        }
93    }
94
95    /// Commits staged changes: releases all pending old roots for the account.
96    pub fn commit_roots(&mut self, account_id: AccountId) {
97        if let Some(old_roots_stack) = self.pending_old_roots.remove(&account_id) {
98            for old_roots in old_roots_stack {
99                let to_pop = decrement_refcounts(&mut self.root_refcounts, &old_roots);
100                self.safe_pop_smts(to_pop);
101            }
102        }
103    }
104
105    /// Discards the most recent staged change: restores old roots and releases new roots.
106    ///
107    /// If there are old roots to restore, the current roots are replaced with them.
108    /// If there are no old roots (i.e., the account was first staged without prior state),
109    /// the current roots are simply removed.
110    pub fn discard_roots(&mut self, account_id: AccountId) {
111        let old_roots = self.pending_old_roots.get_mut(&account_id).and_then(Vec::pop);
112
113        // Release the current (staged) roots and restore old ones if available
114        let new_roots = match old_roots {
115            Some(old_roots) => self.account_roots.insert(account_id, old_roots),
116            None => self.account_roots.remove(&account_id),
117        };
118
119        if let Some(new_roots) = new_roots {
120            let to_pop = decrement_refcounts(&mut self.root_refcounts, &new_roots);
121            self.safe_pop_smts(to_pop);
122        }
123
124        // Clean up empty stack
125        if self.pending_old_roots.get(&account_id).is_some_and(Vec::is_empty) {
126            self.pending_old_roots.remove(&account_id);
127        }
128    }
129
130    /// Replaces roots atomically: sets new roots and immediately releases old roots.
131    ///
132    /// Use this when no rollback is needed (e.g., initial insert, network updates).
133    ///
134    /// # Panics
135    ///
136    /// Panics if there are pending staged changes for the account. Use
137    /// [`Self::commit_roots`] or [`Self::discard_roots`] first.
138    pub fn replace_roots(&mut self, account_id: AccountId, new_roots: Vec<Word>) {
139        assert!(
140            !self.pending_old_roots.contains_key(&account_id),
141            "cannot replace roots while staged changes are pending for account {account_id}"
142        );
143        increment_refcounts(&mut self.root_refcounts, &new_roots);
144        if let Some(old_roots) = self.account_roots.insert(account_id, new_roots) {
145            let to_pop = decrement_refcounts(&mut self.root_refcounts, &old_roots);
146            self.safe_pop_smts(to_pop);
147        }
148    }
149
150    // TREE MUTATORS
151    // --------------------------------------------------------------------------------------------
152
153    /// Updates the SMT forest with the new asset values.
154    pub fn update_asset_nodes(
155        &mut self,
156        root: Word,
157        new_assets: impl Iterator<Item = Asset>,
158        removed_vault_ids: impl Iterator<Item = AssetId>,
159    ) -> Result<Word, StoreError> {
160        let entries: Vec<(Word, Word)> = new_assets
161            .map(|asset| {
162                let key: Word = asset.id().hash().into();
163                let value = asset.to_value_word();
164                (key, value)
165            })
166            .chain(removed_vault_ids.map(|vault_id| (vault_id.hash().into(), EMPTY_WORD)))
167            .collect();
168
169        if entries.is_empty() {
170            return Ok(root);
171        }
172
173        let new_root = self.forest.batch_insert(root, entries).map_err(StoreError::from)?;
174        Ok(new_root)
175    }
176
177    /// Updates the SMT forest with the new storage map values.
178    pub fn update_storage_map_nodes(
179        &mut self,
180        root: Word,
181        entries: impl Iterator<Item = (StorageMapKey, Word)>,
182    ) -> Result<Word, StoreError> {
183        let entries: Vec<(StorageMapKeyHash, Word)> =
184            entries.map(|(key, value)| (key.hash(), value)).collect();
185
186        if entries.is_empty() {
187            return Ok(root);
188        }
189
190        let new_root = self
191            .forest
192            .batch_insert(root, entries.into_iter().map(|(key, value)| (Word::from(key), value)))
193            .map_err(StoreError::from)?;
194        Ok(new_root)
195    }
196
197    /// Inserts the asset vault SMT nodes to the SMT forest.
198    pub fn insert_asset_nodes(&mut self, vault: &AssetVault) -> Result<(), StoreError> {
199        let smt = Smt::with_entries(vault.assets().map(|asset| {
200            let key: Word = asset.id().hash().into();
201            let value = asset.to_value_word();
202            (key, value)
203        }))
204        .map_err(StoreError::from)?;
205
206        let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
207        let entries: Vec<(Word, Word)> = smt.entries().map(|(k, v)| (*k, *v)).collect();
208        if entries.is_empty() {
209            return Ok(());
210        }
211        let new_root = self.forest.batch_insert(empty_root, entries).map_err(StoreError::from)?;
212        debug_assert_eq!(new_root, smt.root());
213        Ok(())
214    }
215
216    /// Inserts all storage map SMT nodes to the SMT forest.
217    pub fn insert_storage_map_nodes(&mut self, storage: &AccountStorage) -> Result<(), StoreError> {
218        let maps = storage.slots().iter().filter_map(|slot| match slot.content() {
219            StorageSlotContent::Map(map) => Some(map),
220            StorageSlotContent::Value(_) => None,
221        });
222
223        for map in maps {
224            self.insert_storage_map_nodes_for_map(map)?;
225        }
226        Ok(())
227    }
228
229    /// Inserts the SMT nodes for an account's vault and storage maps into the
230    /// forest, without tracking roots for the account.
231    pub fn insert_account_state(
232        &mut self,
233        vault: &AssetVault,
234        storage: &AccountStorage,
235    ) -> Result<(), StoreError> {
236        self.insert_storage_map_nodes(storage)?;
237        self.insert_asset_nodes(vault)?;
238        Ok(())
239    }
240
241    /// Inserts all SMT nodes for an account's vault and storage, then stages
242    /// the account's roots for later commit or discard.
243    pub fn insert_and_stage_account_state(
244        &mut self,
245        account_id: AccountId,
246        vault: &AssetVault,
247        storage: &AccountStorage,
248    ) -> Result<(), StoreError> {
249        self.insert_account_state(vault, storage)?;
250        let roots = Self::collect_account_roots(vault, storage);
251        self.stage_roots(account_id, roots);
252        Ok(())
253    }
254
255    /// Inserts all SMT nodes for an account's vault and storage, then replaces
256    /// the account's tracked roots atomically.
257    pub fn insert_and_register_account_state(
258        &mut self,
259        account_id: AccountId,
260        vault: &AssetVault,
261        storage: &AccountStorage,
262    ) -> Result<(), StoreError> {
263        self.insert_account_state(vault, storage)?;
264        let roots = Self::collect_account_roots(vault, storage);
265        self.replace_roots(account_id, roots);
266        Ok(())
267    }
268
269    /// Inserts storage map SMT nodes for a specific storage map.
270    pub fn insert_storage_map_nodes_for_map(&mut self, map: &StorageMap) -> Result<(), StoreError> {
271        let empty_root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
272        let entries: Vec<(StorageMapKeyHash, Word)> =
273            map.entries().map(|(key, value)| (key.hash(), *value)).collect();
274        if entries.is_empty() {
275            return Ok(());
276        }
277        self.forest
278            .batch_insert(
279                empty_root,
280                entries.into_iter().map(|(key, value)| (Word::from(key), value)),
281            )
282            .map_err(StoreError::from)?;
283        Ok(())
284    }
285
286    // HELPERS
287    // --------------------------------------------------------------------------------------------
288
289    /// Collects all SMT roots (vault root + storage map roots) for an account's state.
290    fn collect_account_roots(vault: &AssetVault, storage: &AccountStorage) -> Vec<Word> {
291        let mut roots = vec![vault.root()];
292        for slot in storage.slots() {
293            if let StorageSlotContent::Map(map) = slot.content() {
294                roots.push(map.root());
295            }
296        }
297        roots
298    }
299
300    /// Pops SMT roots from the forest that are no longer referenced by any account.
301    fn safe_pop_smts(&mut self, roots: impl IntoIterator<Item = Word>) {
302        self.forest.pop_smts(roots);
303    }
304}
305
306fn increment_refcounts(refcounts: &mut BTreeMap<Word, usize>, roots: &[Word]) {
307    for root in roots {
308        *refcounts.entry(*root).or_insert(0) += 1;
309    }
310}
311
312/// Decrements refcounts for the given roots, returning those that reached zero.
313fn decrement_refcounts(refcounts: &mut BTreeMap<Word, usize>, roots: &[Word]) -> Vec<Word> {
314    let mut to_pop = Vec::new();
315    for root in roots {
316        if let Some(count) = refcounts.get_mut(root) {
317            *count -= 1;
318            if *count == 0 {
319                refcounts.remove(root);
320                to_pop.push(*root);
321            }
322        }
323    }
324    to_pop
325}
326
327#[cfg(test)]
328mod tests {
329    use miden_protocol::testing::account_id::{
330        ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
331        ACCOUNT_ID_PUBLIC_NON_FUNGIBLE_FAUCET,
332    };
333    use miden_protocol::{ONE, ZERO};
334
335    use super::*;
336
337    fn account_a() -> AccountId {
338        AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap()
339    }
340
341    fn account_b() -> AccountId {
342        AccountId::try_from(ACCOUNT_ID_PUBLIC_NON_FUNGIBLE_FAUCET).unwrap()
343    }
344
345    /// Creates a `StorageMap` with a single entry and inserts its nodes into the forest.
346    /// Returns the map's root.
347    fn insert_map(forest: &mut AccountSmtForest, key: Word, value: Word) -> Word {
348        let mut map = StorageMap::new();
349        map.insert(StorageMapKey::new(key), value).unwrap();
350        forest.insert_storage_map_nodes_for_map(&map).unwrap();
351        map.root()
352    }
353
354    /// Returns true if the forest can still serve a proof for the given root.
355    fn root_is_live(forest: &AccountSmtForest, root: Word, key: Word) -> bool {
356        forest.get_storage_map_item_witness(root, StorageMapKey::new(key)).is_ok()
357    }
358
359    #[test]
360    fn stage_then_commit_releases_old_roots() {
361        let mut forest = AccountSmtForest::new();
362        let id = account_a();
363
364        let key1: Word = [ONE, ZERO, ZERO, ZERO].into();
365        let key2: Word = [ZERO, ONE, ZERO, ZERO].into();
366        let val: Word = [ONE, ONE, ONE, ONE].into();
367
368        let root1 = insert_map(&mut forest, key1, val);
369        let root2 = insert_map(&mut forest, key2, val);
370
371        // Initial state
372        forest.replace_roots(id, vec![root1]);
373        assert_eq!(forest.get_roots(&id), Some(&vec![root1]));
374
375        // Stage new roots (apply_delta)
376        forest.stage_roots(id, vec![root2]);
377        assert_eq!(forest.get_roots(&id), Some(&vec![root2]));
378
379        // Both roots alive during staging (old preserved for rollback)
380        assert!(root_is_live(&forest, root1, key1));
381        assert!(root_is_live(&forest, root2, key2));
382
383        // Commit — old roots released
384        forest.commit_roots(id);
385        assert_eq!(forest.get_roots(&id), Some(&vec![root2]));
386        assert!(!root_is_live(&forest, root1, key1));
387        assert!(root_is_live(&forest, root2, key2));
388    }
389
390    #[test]
391    fn stage_then_discard_restores_old_roots() {
392        let mut forest = AccountSmtForest::new();
393        let id = account_a();
394
395        let key1: Word = [ONE, ZERO, ZERO, ZERO].into();
396        let key2: Word = [ZERO, ONE, ZERO, ZERO].into();
397        let val: Word = [ONE, ONE, ONE, ONE].into();
398
399        let root1 = insert_map(&mut forest, key1, val);
400        let root2 = insert_map(&mut forest, key2, val);
401
402        forest.replace_roots(id, vec![root1]);
403
404        // Stage and discard (rollback)
405        forest.stage_roots(id, vec![root2]);
406        forest.discard_roots(id);
407
408        assert_eq!(forest.get_roots(&id), Some(&vec![root1]));
409        assert!(root_is_live(&forest, root1, key1));
410        assert!(!root_is_live(&forest, root2, key2));
411    }
412
413    #[test]
414    fn shared_root_survives_single_account_replacement() {
415        let mut forest = AccountSmtForest::new();
416        let id1 = account_a();
417        let id2 = account_b();
418
419        let key: Word = [ONE, ZERO, ZERO, ZERO].into();
420        let val: Word = [ONE, ONE, ONE, ONE].into();
421        let shared_root = insert_map(&mut forest, key, val);
422
423        // Both accounts reference the same root
424        forest.replace_roots(id1, vec![shared_root]);
425        forest.replace_roots(id2, vec![shared_root]);
426
427        // Replace id1 with a different root
428        let key2: Word = [ZERO, ONE, ZERO, ZERO].into();
429        let other_root = insert_map(&mut forest, key2, val);
430        forest.replace_roots(id1, vec![other_root]);
431
432        // Shared root still alive (id2 still references it)
433        assert!(root_is_live(&forest, shared_root, key));
434
435        // Replace id2 too — now shared root should be popped
436        forest.replace_roots(id2, vec![other_root]);
437        assert!(!root_is_live(&forest, shared_root, key));
438    }
439
440    #[test]
441    fn multiple_stages_discard_one_at_a_time() {
442        let mut forest = AccountSmtForest::new();
443        let id = account_a();
444
445        let key_a: Word = [ONE, ZERO, ZERO, ZERO].into();
446        let key_b: Word = [ZERO, ONE, ZERO, ZERO].into();
447        let key_c: Word = [ZERO, ZERO, ONE, ZERO].into();
448        let val: Word = [ONE, ONE, ONE, ONE].into();
449
450        let root_a = insert_map(&mut forest, key_a, val);
451        let root_b = insert_map(&mut forest, key_b, val);
452        let root_c = insert_map(&mut forest, key_c, val);
453
454        // A -> B -> C
455        forest.replace_roots(id, vec![root_a]);
456        forest.stage_roots(id, vec![root_b]);
457        forest.stage_roots(id, vec![root_c]);
458        assert_eq!(forest.get_roots(&id), Some(&vec![root_c]));
459
460        // Discard C -> back to B
461        forest.discard_roots(id);
462        assert_eq!(forest.get_roots(&id), Some(&vec![root_b]));
463        assert!(!root_is_live(&forest, root_c, key_c));
464        assert!(root_is_live(&forest, root_b, key_b));
465        assert!(root_is_live(&forest, root_a, key_a));
466
467        // Discard B -> back to A
468        forest.discard_roots(id);
469        assert_eq!(forest.get_roots(&id), Some(&vec![root_a]));
470        assert!(!root_is_live(&forest, root_b, key_b));
471        assert!(root_is_live(&forest, root_a, key_a));
472    }
473
474    #[test]
475    fn multiple_stages_commit_releases_all_old() {
476        let mut forest = AccountSmtForest::new();
477        let id = account_a();
478
479        let key_a: Word = [ONE, ZERO, ZERO, ZERO].into();
480        let key_b: Word = [ZERO, ONE, ZERO, ZERO].into();
481        let key_c: Word = [ZERO, ZERO, ONE, ZERO].into();
482        let val: Word = [ONE, ONE, ONE, ONE].into();
483
484        let root_a = insert_map(&mut forest, key_a, val);
485        let root_b = insert_map(&mut forest, key_b, val);
486        let root_c = insert_map(&mut forest, key_c, val);
487
488        // A -> B -> C, then commit
489        forest.replace_roots(id, vec![root_a]);
490        forest.stage_roots(id, vec![root_b]);
491        forest.stage_roots(id, vec![root_c]);
492        forest.commit_roots(id);
493
494        // Only C survives
495        assert_eq!(forest.get_roots(&id), Some(&vec![root_c]));
496        assert!(!root_is_live(&forest, root_a, key_a));
497        assert!(!root_is_live(&forest, root_b, key_b));
498        assert!(root_is_live(&forest, root_c, key_c));
499    }
500
501    #[test]
502    fn unchanged_root_survives_stage_commit() {
503        let mut forest = AccountSmtForest::new();
504        let id = account_a();
505
506        let key1: Word = [ONE, ZERO, ZERO, ZERO].into();
507        let key2: Word = [ZERO, ONE, ZERO, ZERO].into();
508        let val: Word = [ONE, ONE, ONE, ONE].into();
509
510        let shared_root = insert_map(&mut forest, key1, val);
511        let changing_root = insert_map(&mut forest, key2, val);
512
513        // Initial: [shared, changing]
514        forest.replace_roots(id, vec![shared_root, changing_root]);
515
516        // Delta only changes the second root; shared_root stays
517        let key3: Word = [ZERO, ZERO, ONE, ZERO].into();
518        let new_root = insert_map(&mut forest, key3, val);
519        forest.stage_roots(id, vec![shared_root, new_root]);
520        forest.commit_roots(id);
521
522        // shared_root must survive (it's in both old and new)
523        assert!(root_is_live(&forest, shared_root, key1));
524        // changing_root should be popped
525        assert!(!root_is_live(&forest, changing_root, key2));
526        // new_root should be alive
527        assert!(root_is_live(&forest, new_root, key3));
528    }
529}
miden_client/store/smt_forest.rs

miden_client/store/
smt_forest.rs
