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

miden_client/store/
smt_forest.rs
