miden_objects/block/
account_tree.rs

1use alloc::string::ToString;
2use alloc::vec::Vec;
3
4use miden_core::utils::{ByteReader, ByteWriter, Deserializable, Serializable};
5use miden_crypto::merkle::{MerkleError, MutationSet, Smt, SmtLeaf};
6use miden_processor::{DeserializationError, SMT_DEPTH};
7
8use crate::account::{AccountId, AccountIdPrefix};
9use crate::block::AccountWitness;
10use crate::errors::AccountTreeError;
11use crate::{Felt, Word};
12
13// ACCOUNT TREE
14// ================================================================================================
15
16/// The sparse merkle tree of all accounts in the blockchain.
17///
18/// The key is the [`AccountId`] while the value is the current state commitment of the account,
19/// i.e. [`Account::commitment`](crate::account::Account::commitment). If the account is new, then
20/// the commitment is the [`EMPTY_WORD`](crate::EMPTY_WORD).
21///
22/// Each account ID occupies exactly one leaf in the tree, which is identified by its
23/// [`AccountId::prefix`]. In other words, account ID prefixes are unique in the blockchain.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct AccountTree {
26    smt: Smt,
27}
28
29impl AccountTree {
30    // CONSTANTS
31    // --------------------------------------------------------------------------------------------
32
33    /// The depth of the account tree.
34    pub const DEPTH: u8 = SMT_DEPTH;
35
36    /// The index of the account ID suffix in the SMT key.
37    pub(super) const KEY_SUFFIX_IDX: usize = 2;
38    /// The index of the account ID prefix in the SMT key.
39    pub(super) const KEY_PREFIX_IDX: usize = 3;
40
41    // CONSTRUCTORS
42    // --------------------------------------------------------------------------------------------
43
44    /// Creates a new, empty account tree.
45    pub fn new() -> Self {
46        AccountTree { smt: Smt::new() }
47    }
48
49    /// Returns a new [`Smt`] instantiated with the provided entries.
50    ///
51    /// If the `concurrent` feature of `miden-crypto` is enabled, this function uses a parallel
52    /// implementation to process the entries efficiently, otherwise it defaults to the
53    /// sequential implementation.
54    ///
55    /// # Errors
56    ///
57    /// Returns an error if:
58    /// - the provided entries contain multiple commitments for the same account ID.
59    /// - multiple account IDs share the same prefix.
60    pub fn with_entries<I>(
61        entries: impl IntoIterator<Item = (AccountId, Word), IntoIter = I>,
62    ) -> Result<Self, AccountTreeError>
63    where
64        I: ExactSizeIterator<Item = (AccountId, Word)>,
65    {
66        let entries = entries.into_iter();
67        let num_accounts = entries.len();
68
69        let smt = Smt::with_entries(
70            entries.map(|(id, commitment)| (Self::id_to_smt_key(id), commitment)),
71        )
72        .map_err(|err| {
73            let MerkleError::DuplicateValuesForIndex(leaf_idx) = err else {
74                unreachable!("the only error returned by Smt::with_entries is of this type");
75            };
76
77            // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed that
78            // the leaf_idx is a valid Felt as well as a valid account ID prefix.
79            AccountTreeError::DuplicateStateCommitments {
80                prefix: AccountIdPrefix::new_unchecked(
81                    Felt::try_from(leaf_idx).expect("leaf index should be a valid felt"),
82                ),
83            }
84        })?;
85
86        // If the number of leaves in the SMT is smaller than the number of accounts that were
87        // passed in, it means that at least one account ID pair ended up in the same leaf. If this
88        // is the case, we iterate the SMT entries to find the duplicated account ID prefix.
89        if smt.num_leaves() < num_accounts {
90            for (leaf_idx, leaf) in smt.leaves() {
91                if leaf.num_entries() >= 2 {
92                    // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed
93                    // that the leaf_idx is a valid Felt as well as a valid
94                    // account ID prefix.
95                    return Err(AccountTreeError::DuplicateIdPrefix {
96                        duplicate_prefix: AccountIdPrefix::new_unchecked(
97                            Felt::try_from(leaf_idx.value())
98                                .expect("leaf index should be a valid felt"),
99                        ),
100                    });
101                }
102            }
103        }
104
105        Ok(AccountTree { smt })
106    }
107
108    // PUBLIC ACCESSORS
109    // --------------------------------------------------------------------------------------------
110
111    /// Returns an opening of the leaf associated with the `account_id`. This is a proof of the
112    /// current state commitment of the given account ID.
113    ///
114    /// Conceptually, an opening is a Merkle path to the leaf, as well as the leaf itself.
115    pub fn open(&self, account_id: AccountId) -> AccountWitness {
116        let key = Self::id_to_smt_key(account_id);
117        let proof = self.smt.open(&key);
118
119        AccountWitness::from_smt_proof(account_id, proof)
120    }
121
122    /// Returns the current state commitment of the given account ID.
123    pub fn get(&self, account_id: AccountId) -> Word {
124        let key = Self::id_to_smt_key(account_id);
125        self.smt.get_value(&key)
126    }
127
128    /// Returns the root of the tree.
129    pub fn root(&self) -> Word {
130        self.smt.root()
131    }
132
133    /// Returns true if the tree contains a leaf for the given account ID prefix.
134    pub fn contains_account_id_prefix(&self, account_id_prefix: AccountIdPrefix) -> bool {
135        let key = Self::id_prefix_to_smt_key(account_id_prefix);
136        let is_empty = matches!(self.smt.get_leaf(&key), SmtLeaf::Empty(_));
137        !is_empty
138    }
139
140    /// Returns the number of account IDs in this tree.
141    pub fn num_accounts(&self) -> usize {
142        // Because each ID's prefix is unique in the tree and occupies a single leaf, the number of
143        // IDs in the tree is equivalent to the number of leaves in the tree.
144        self.smt.num_leaves()
145    }
146
147    /// Returns an iterator over the account ID state commitment pairs in the tree.
148    pub fn account_commitments(&self) -> impl Iterator<Item = (AccountId, Word)> {
149        self.smt.leaves().map(|(_leaf_idx, leaf)| {
150            // SAFETY: By construction no Multiple variant is ever present in the tree.
151            // The Empty variant is not returned by Smt::leaves, because it only returns leaves that
152            // are actually present.
153            let SmtLeaf::Single((key, commitment)) = leaf else {
154                unreachable!("empty and multiple variant should never be encountered")
155            };
156
157            (
158                // SAFETY: By construction, the tree only contains valid IDs.
159                AccountId::try_from([key[Self::KEY_PREFIX_IDX], key[Self::KEY_SUFFIX_IDX]])
160                    .expect("account tree should only contain valid IDs"),
161                *commitment,
162            )
163        })
164    }
165
166    /// Computes the necessary changes to insert the specified (account ID, state commitment) pairs
167    /// into this tree, allowing for validation before applying those changes.
168    ///
169    /// [`Self::apply_mutations`] can be used in order to commit these changes to the tree.
170    ///
171    /// If the `concurrent` feature of `miden-crypto` is enabled, this function uses a parallel
172    /// implementation to compute the mutations, otherwise it defaults to the sequential
173    /// implementation.
174    ///
175    /// This is a thin wrapper around [`Smt::compute_mutations`]. See its documentation for more
176    /// details.
177    ///
178    /// # Errors
179    ///
180    /// Returns an error if:
181    /// - an insertion of an account ID would violate the uniqueness of account ID prefixes in the
182    ///   tree.
183    pub fn compute_mutations(
184        &self,
185        account_commitments: impl IntoIterator<Item = (AccountId, Word)>,
186    ) -> Result<AccountMutationSet, AccountTreeError> {
187        let mutation_set = self.smt.compute_mutations(
188            account_commitments
189                .into_iter()
190                .map(|(id, commitment)| (Self::id_to_smt_key(id), commitment)),
191        );
192
193        for id_key in mutation_set.new_pairs().keys() {
194            // Check if the insertion would be valid.
195            match self.smt.get_leaf(id_key) {
196                // Inserting into an empty leaf is valid.
197                SmtLeaf::Empty(_) => (),
198                SmtLeaf::Single((existing_key, _)) => {
199                    // If the key matches the existing one, then we're updating the leaf, which is
200                    // valid. If it does not match, then we would insert a duplicate.
201                    if existing_key != *id_key {
202                        return Err(AccountTreeError::DuplicateIdPrefix {
203                            duplicate_prefix: Self::smt_key_to_id(*id_key).prefix(),
204                        });
205                    }
206                },
207                SmtLeaf::Multiple(_) => {
208                    unreachable!(
209                        "account tree should never contain duplicate ID prefixes and therefore never a multiple leaf"
210                    )
211                },
212            }
213        }
214
215        Ok(AccountMutationSet::new(mutation_set))
216    }
217
218    // PUBLIC MUTATORS
219    // --------------------------------------------------------------------------------------------
220
221    /// Inserts the state commitment for the given account ID, returning the previous state
222    /// commitment associated with that ID.
223    ///
224    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
225    /// updating the root itself.
226    ///
227    /// # Errors
228    ///
229    /// Returns an error if:
230    /// - the prefix of the account ID already exists in the tree.
231    pub fn insert(
232        &mut self,
233        account_id: AccountId,
234        state_commitment: Word,
235    ) -> Result<Word, AccountTreeError> {
236        let key = Self::id_to_smt_key(account_id);
237        let prev_value = self.smt.insert(key, state_commitment);
238
239        // If the leaf of the account ID now has two or more entries, we've inserted a duplicate
240        // prefix.
241        if self.smt.get_leaf(&key).num_entries() >= 2 {
242            return Err(AccountTreeError::DuplicateIdPrefix {
243                duplicate_prefix: account_id.prefix(),
244            });
245        }
246
247        Ok(prev_value)
248    }
249
250    /// Applies the prospective mutations computed with [`Self::compute_mutations`] to this tree.
251    ///
252    /// # Errors
253    ///
254    /// Returns an error if:
255    /// - `mutations` was computed on a tree with a different root than this one.
256    pub fn apply_mutations(
257        &mut self,
258        mutations: AccountMutationSet,
259    ) -> Result<(), AccountTreeError> {
260        self.smt
261            .apply_mutations(mutations.into_mutation_set())
262            .map_err(AccountTreeError::ApplyMutations)
263    }
264
265    // HELPERS
266    // --------------------------------------------------------------------------------------------
267
268    /// Returns the SMT key of the given account ID.
269    pub(super) fn id_to_smt_key(account_id: AccountId) -> Word {
270        // We construct this in such a way that we're forced to use the constants, so that when
271        // they're updated, the other usages of the constants are also updated.
272        let mut key = Word::empty();
273        key[Self::KEY_SUFFIX_IDX] = account_id.suffix();
274        key[Self::KEY_PREFIX_IDX] = account_id.prefix().as_felt();
275
276        key
277    }
278
279    /// Returns the SMT key of the given account ID prefix.
280    fn id_prefix_to_smt_key(account_id: AccountIdPrefix) -> Word {
281        // We construct this in such a way that we're forced to use the constants, so that when
282        // they're updated, the other usages of the constants are also updated.
283        let mut key = Word::empty();
284        key[Self::KEY_PREFIX_IDX] = account_id.as_felt();
285
286        key
287    }
288
289    /// Returns the [`AccountId`] recovered from the given SMT key.
290    ///
291    /// # Panics
292    ///
293    /// Panics if:
294    /// - the key is not a valid account ID. This should not happen when used on keys from (partial)
295    ///   account tree.
296    pub(super) fn smt_key_to_id(key: Word) -> AccountId {
297        AccountId::try_from([key[Self::KEY_PREFIX_IDX], key[Self::KEY_SUFFIX_IDX]])
298            .expect("account tree should only contain valid IDs")
299    }
300}
301
302impl Default for AccountTree {
303    fn default() -> Self {
304        Self::new()
305    }
306}
307
308// SERIALIZATION
309// ================================================================================================
310
311impl Serializable for AccountTree {
312    fn write_into<W: ByteWriter>(&self, target: &mut W) {
313        self.account_commitments().collect::<Vec<_>>().write_into(target);
314    }
315}
316
317impl Deserializable for AccountTree {
318    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
319        let entries = Vec::<(AccountId, Word)>::read_from(source)?;
320        Self::with_entries(entries)
321            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
322    }
323}
324
325// ACCOUNT MUTATION SET
326// ================================================================================================
327
328/// A newtype wrapper around a [`MutationSet`] for use in the [`AccountTree`].
329///
330/// It guarantees that applying the contained mutations will result in an account tree with unique
331/// account ID prefixes.
332///
333/// It is returned by and used in methods on the [`AccountTree`].
334#[derive(Debug, Clone, PartialEq, Eq)]
335pub struct AccountMutationSet {
336    mutation_set: MutationSet<{ AccountTree::DEPTH }, Word, Word>,
337}
338
339impl AccountMutationSet {
340    // CONSTRUCTORS
341    // --------------------------------------------------------------------------------------------
342
343    /// Creates a new [`AccountMutationSet`] from the provided raw mutation set.
344    fn new(mutation_set: MutationSet<{ AccountTree::DEPTH }, Word, Word>) -> Self {
345        Self { mutation_set }
346    }
347
348    // PUBLIC ACCESSORS
349    // --------------------------------------------------------------------------------------------
350
351    /// Returns a reference to the underlying [`MutationSet`].
352    pub fn as_mutation_set(&self) -> &MutationSet<{ AccountTree::DEPTH }, Word, Word> {
353        &self.mutation_set
354    }
355
356    // PUBLIC MUTATORS
357    // --------------------------------------------------------------------------------------------
358
359    /// Consumes self and returns the underlying [`MutationSet`].
360    pub fn into_mutation_set(self) -> MutationSet<{ AccountTree::DEPTH }, Word, Word> {
361        self.mutation_set
362    }
363}
364
365// TESTS
366// ================================================================================================
367
368#[cfg(test)]
369pub(super) mod tests {
370    use std::vec::Vec;
371
372    use assert_matches::assert_matches;
373
374    use super::*;
375    use crate::account::{AccountStorageMode, AccountType};
376    use crate::testing::account_id::{AccountIdBuilder, account_id};
377
378    pub(crate) fn setup_duplicate_prefix_ids() -> [(AccountId, Word); 2] {
379        let id0 = AccountId::try_from(account_id(
380            AccountType::FungibleFaucet,
381            AccountStorageMode::Public,
382            0xaabb_ccdd,
383        ))
384        .unwrap();
385        let id1 = AccountId::try_from(account_id(
386            AccountType::FungibleFaucet,
387            AccountStorageMode::Public,
388            0xaabb_ccff,
389        ))
390        .unwrap();
391        assert_eq!(id0.prefix(), id1.prefix(), "test requires that these ids have the same prefix");
392
393        let commitment0 = Word::from([0, 0, 0, 42u32]);
394        let commitment1 = Word::from([0, 0, 0, 24u32]);
395
396        assert_eq!(id0.prefix(), id1.prefix(), "test requires that these ids have the same prefix");
397        [(id0, commitment0), (id1, commitment1)]
398    }
399
400    #[test]
401    fn insert_fails_on_duplicate_prefix() {
402        let mut tree = AccountTree::new();
403        let [(id0, commitment0), (id1, commitment1)] = setup_duplicate_prefix_ids();
404
405        tree.insert(id0, commitment0).unwrap();
406        assert_eq!(tree.get(id0), commitment0);
407
408        let err = tree.insert(id1, commitment1).unwrap_err();
409
410        assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
411          duplicate_prefix
412        } if duplicate_prefix == id0.prefix());
413    }
414
415    #[test]
416    fn with_entries_fails_on_duplicate_prefix() {
417        let entries = setup_duplicate_prefix_ids();
418
419        let err = AccountTree::with_entries(entries.iter().copied()).unwrap_err();
420
421        assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
422          duplicate_prefix
423        } if duplicate_prefix == entries[0].0.prefix());
424    }
425
426    #[test]
427    fn insert_succeeds_on_multiple_updates() {
428        let mut tree = AccountTree::new();
429        let [(id0, commitment0), (_, commitment1)] = setup_duplicate_prefix_ids();
430
431        tree.insert(id0, commitment0).unwrap();
432        tree.insert(id0, commitment1).unwrap();
433        assert_eq!(tree.get(id0), commitment1);
434    }
435
436    #[test]
437    fn apply_mutations() {
438        let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
439        let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
440        let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
441
442        let digest0 = Word::from([0, 0, 0, 1u32]);
443        let digest1 = Word::from([0, 0, 0, 2u32]);
444        let digest2 = Word::from([0, 0, 0, 3u32]);
445        let digest3 = Word::from([0, 0, 0, 4u32]);
446
447        let mut tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap();
448
449        let mutations = tree
450            .compute_mutations([(id0, digest1), (id1, digest2), (id2, digest3)])
451            .unwrap();
452
453        tree.apply_mutations(mutations).unwrap();
454
455        assert_eq!(tree.num_accounts(), 3);
456        assert_eq!(tree.get(id0), digest1);
457        assert_eq!(tree.get(id1), digest2);
458        assert_eq!(tree.get(id2), digest3);
459    }
460
461    #[test]
462    fn duplicates_in_compute_mutations() {
463        let [pair0, pair1] = setup_duplicate_prefix_ids();
464        let id2 = AccountIdBuilder::new().build_with_seed([5; 32]);
465        let commitment2 = Word::from([0, 0, 0, 99u32]);
466
467        let tree = AccountTree::with_entries([pair0, (id2, commitment2)]).unwrap();
468
469        let err = tree.compute_mutations([pair1]).unwrap_err();
470
471        assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
472          duplicate_prefix
473        } if duplicate_prefix == pair1.0.prefix());
474    }
475
476    #[test]
477    fn account_commitments() {
478        let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
479        let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
480        let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
481
482        let digest0 = Word::from([0, 0, 0, 1u32]);
483        let digest1 = Word::from([0, 0, 0, 2u32]);
484        let digest2 = Word::from([0, 0, 0, 3u32]);
485        let empty_digest = Word::empty();
486
487        let mut tree =
488            AccountTree::with_entries([(id0, digest0), (id1, digest1), (id2, digest2)]).unwrap();
489
490        // remove id2
491        tree.insert(id2, empty_digest).unwrap();
492
493        assert_eq!(tree.num_accounts(), 2);
494
495        let accounts: Vec<_> = tree.account_commitments().collect();
496        assert_eq!(accounts.len(), 2);
497        assert!(accounts.contains(&(id0, digest0)));
498        assert!(accounts.contains(&(id1, digest1)));
499    }
500
501    #[test]
502    fn account_witness() {
503        let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
504        let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
505
506        let digest0 = Word::from([0, 0, 0, 1u32]);
507        let digest1 = Word::from([0, 0, 0, 2u32]);
508
509        let tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap();
510
511        assert_eq!(tree.num_accounts(), 2);
512
513        for id in [id0, id1] {
514            let (control_path, control_leaf) =
515                tree.smt.open(&AccountTree::id_to_smt_key(id)).into_parts();
516            let witness = tree.open(id);
517
518            assert_eq!(witness.leaf(), control_leaf);
519            assert_eq!(witness.path(), &control_path);
520        }
521    }
522
523    #[test]
524    fn contains_account_prefix() {
525        // Create a tree with a single account.
526        let [pair0, pair1] = setup_duplicate_prefix_ids();
527        let tree = AccountTree::with_entries([(pair0.0, pair0.1)]).unwrap();
528        assert_eq!(tree.num_accounts(), 1);
529
530        // Validate the leaf for the inserted account exists.
531        assert!(tree.contains_account_id_prefix(pair0.0.prefix()));
532
533        // Validate the leaf for the uninserted account with same prefix exists.
534        assert!(tree.contains_account_id_prefix(pair1.0.prefix()));
535
536        // Validate the unrelated, uninserted account leaf does not exist.
537        let id1 = AccountIdBuilder::new().build_with_seed([7; 32]);
538        assert!(!tree.contains_account_id_prefix(id1.prefix()));
539    }
540}
miden_objects/block/account_tree.rs

miden_objects/block/
account_tree.rs
