1use alloc::string::ToString;
2use alloc::vec::Vec;
3
4use crate::Word;
5use crate::account::{AccountId, AccountIdPrefix};
6use crate::crypto::merkle::MerkleError;
7use crate::crypto::merkle::smt::{MutationSet, SMT_DEPTH, Smt, SmtLeaf};
8use crate::errors::AccountTreeError;
9use crate::utils::serde::{
10 ByteReader,
11 ByteWriter,
12 Deserializable,
13 DeserializationError,
14 Serializable,
15};
16
17mod partial;
18pub use partial::PartialAccountTree;
19
20mod witness;
21pub use witness::AccountWitness;
22
23mod backend;
24pub use backend::AccountTreeBackend;
25
26mod account_id_key;
27pub use account_id_key::AccountIdKey;
28
29#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct AccountTree<S = Smt> {
42 smt: S,
43}
44
45impl<S> Default for AccountTree<S>
46where
47 S: Default,
48{
49 fn default() -> Self {
50 Self { smt: Default::default() }
51 }
52}
53
54impl<S> AccountTree<S>
55where
56 S: AccountTreeBackend<Error = MerkleError>,
57{
58 pub const DEPTH: u8 = SMT_DEPTH;
63
64 pub(super) const KEY_SUFFIX_IDX: usize = 2;
66 pub(super) const KEY_PREFIX_IDX: usize = 3;
68
69 pub fn new(smt: S) -> Result<Self, AccountTreeError> {
83 for (_leaf_idx, leaf) in smt.leaves() {
84 match leaf {
85 SmtLeaf::Empty(_) => {
86 continue;
88 },
89 SmtLeaf::Single((key, _)) => {
90 AccountIdKey::try_from_word(key).map_err(|err| {
92 AccountTreeError::InvalidAccountIdKey { key, source: err }
93 })?;
94 },
95 SmtLeaf::Multiple(entries) => {
96 if let Some((key, _)) = entries.first() {
99 let key = *key;
100 let account_id = AccountIdKey::try_from_word(key).map_err(|err| {
101 AccountTreeError::InvalidAccountIdKey { key, source: err }
102 })?;
103
104 return Err(AccountTreeError::DuplicateIdPrefix {
105 duplicate_prefix: account_id.prefix(),
106 });
107 }
108 },
109 }
110 }
111
112 Ok(Self::new_unchecked(smt))
113 }
114
115 pub fn new_unchecked(smt: S) -> Self {
126 AccountTree { smt }
127 }
128
129 pub fn open(&self, account_id: AccountId) -> AccountWitness {
141 let key = AccountIdKey::from(account_id).as_word();
142 let proof = self.smt.open(&key);
143
144 AccountWitness::from_smt_proof(account_id, proof)
145 }
146
147 pub fn get(&self, account_id: AccountId) -> Word {
149 let key = AccountIdKey::from(account_id).as_word();
150 self.smt.get_value(&key)
151 }
152
153 pub fn root(&self) -> Word {
155 self.smt.root()
156 }
157
158 pub fn contains_account_id_prefix(&self, account_id_prefix: AccountIdPrefix) -> bool {
160 let key = Self::id_prefix_to_smt_key(account_id_prefix);
161 let is_empty = matches!(self.smt.get_leaf(&key), SmtLeaf::Empty(_));
162 !is_empty
163 }
164
165 pub fn num_accounts(&self) -> usize {
167 self.smt.num_leaves()
170 }
171
172 pub fn account_commitments(&self) -> impl Iterator<Item = (AccountId, Word)> {
174 self.smt.leaves().map(|(_leaf_idx, leaf)| {
175 let SmtLeaf::Single((key, commitment)) = leaf else {
179 unreachable!("empty and multiple variant should never be encountered")
180 };
181
182 (
183 AccountId::try_from_elements(key[Self::KEY_SUFFIX_IDX], key[Self::KEY_PREFIX_IDX])
185 .expect("account tree should only contain valid IDs"),
186 commitment,
187 )
188 })
189 }
190
191 pub fn compute_mutations(
209 &self,
210 account_commitments: impl IntoIterator<Item = (AccountId, Word)>,
211 ) -> Result<AccountMutationSet, AccountTreeError> {
212 let mutation_set = self
213 .smt
214 .compute_mutations(Vec::from_iter(
215 account_commitments
216 .into_iter()
217 .map(|(id, commitment)| (AccountIdKey::from(id).as_word(), commitment)),
218 ))
219 .map_err(AccountTreeError::ComputeMutations)?;
220
221 for id_key in mutation_set.new_pairs().keys() {
222 match self.smt.get_leaf(id_key) {
224 SmtLeaf::Empty(_) => (),
226 SmtLeaf::Single((existing_key, _)) => {
227 if existing_key != *id_key {
230 return Err(AccountTreeError::DuplicateIdPrefix {
231 duplicate_prefix: AccountIdKey::try_from_word(*id_key)
232 .expect("account tree should only contain valid IDs")
233 .prefix(),
234 });
235 }
236 },
237 SmtLeaf::Multiple(_) => {
238 unreachable!(
239 "account tree should never contain duplicate ID prefixes and therefore never a multiple leaf"
240 )
241 },
242 }
243 }
244
245 Ok(AccountMutationSet::new(mutation_set))
246 }
247
248 pub fn insert(
262 &mut self,
263 account_id: AccountId,
264 state_commitment: Word,
265 ) -> Result<Word, AccountTreeError> {
266 let key = AccountIdKey::from(account_id).as_word();
267 let prev_value = self.smt.insert(key, state_commitment)
270 .expect("account tree should always have a single value per key, and hence cannot exceed the maximum leaf number");
271
272 if self.smt.get_leaf(&key).num_entries() >= 2 {
275 return Err(AccountTreeError::DuplicateIdPrefix {
276 duplicate_prefix: account_id.prefix(),
277 });
278 }
279
280 Ok(prev_value)
281 }
282
283 pub fn apply_mutations(
290 &mut self,
291 mutations: AccountMutationSet,
292 ) -> Result<(), AccountTreeError> {
293 self.smt
294 .apply_mutations(mutations.into_mutation_set())
295 .map_err(AccountTreeError::ApplyMutations)
296 }
297
298 pub fn apply_mutations_with_reversion(
308 &mut self,
309 mutations: AccountMutationSet,
310 ) -> Result<AccountMutationSet, AccountTreeError> {
311 let reversion = self
312 .smt
313 .apply_mutations_with_reversion(mutations.into_mutation_set())
314 .map_err(AccountTreeError::ApplyMutations)?;
315 Ok(AccountMutationSet::new(reversion))
316 }
317
318 fn id_prefix_to_smt_key(account_id: AccountIdPrefix) -> Word {
323 let mut key = Word::empty();
326 key[Self::KEY_PREFIX_IDX] = account_id.as_felt();
327
328 key
329 }
330}
331
332impl Serializable for AccountTree {
336 fn write_into<W: ByteWriter>(&self, target: &mut W) {
337 self.account_commitments().collect::<Vec<_>>().write_into(target);
338 }
339}
340
341impl Deserializable for AccountTree {
342 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
343 let entries = Vec::<(AccountId, Word)>::read_from(source)?;
344
345 let mut seen_prefixes = alloc::collections::BTreeSet::new();
347 for (id, _) in &entries {
348 if !seen_prefixes.insert(id.prefix()) {
349 return Err(DeserializationError::InvalidValue(format!(
350 "Duplicate account ID prefix: {}",
351 id.prefix()
352 )));
353 }
354 }
355
356 let smt = Smt::with_entries(
358 entries.into_iter().map(|(k, v)| (AccountIdKey::from(k).as_word(), v)),
359 )
360 .map_err(|err| DeserializationError::InvalidValue(err.to_string()))?;
361 Ok(Self::new_unchecked(smt))
362 }
363}
364
365#[derive(Debug, Clone, PartialEq, Eq)]
375pub struct AccountMutationSet {
376 mutation_set: MutationSet<SMT_DEPTH, Word, Word>,
377}
378
379impl AccountMutationSet {
380 fn new(mutation_set: MutationSet<SMT_DEPTH, Word, Word>) -> Self {
385 Self { mutation_set }
386 }
387
388 pub fn as_mutation_set(&self) -> &MutationSet<SMT_DEPTH, Word, Word> {
393 &self.mutation_set
394 }
395
396 pub fn into_mutation_set(self) -> MutationSet<SMT_DEPTH, Word, Word> {
401 self.mutation_set
402 }
403}
404
405#[cfg(test)]
409pub(super) mod tests {
410 use std::vec::Vec;
411
412 use assert_matches::assert_matches;
413
414 use super::*;
415 use crate::account::{AccountStorageMode, AccountType};
416 use crate::testing::account_id::{AccountIdBuilder, account_id};
417
418 pub(crate) fn setup_duplicate_prefix_ids() -> [(AccountId, Word); 2] {
419 let id0 = AccountId::try_from(account_id(
420 AccountType::FungibleFaucet,
421 AccountStorageMode::Public,
422 0xaabb_ccdd,
423 ))
424 .unwrap();
425 let id1 = AccountId::try_from(account_id(
426 AccountType::FungibleFaucet,
427 AccountStorageMode::Public,
428 0xaabb_ccff,
429 ))
430 .unwrap();
431 assert_eq!(id0.prefix(), id1.prefix(), "test requires that these ids have the same prefix");
432
433 let commitment0 = Word::from([0, 0, 0, 42u32]);
434 let commitment1 = Word::from([0, 0, 0, 24u32]);
435
436 assert_eq!(id0.prefix(), id1.prefix(), "test requires that these ids have the same prefix");
437 [(id0, commitment0), (id1, commitment1)]
438 }
439
440 #[test]
441 fn insert_fails_on_duplicate_prefix() {
442 let mut tree = AccountTree::<Smt>::default();
443 let [(id0, commitment0), (id1, commitment1)] = setup_duplicate_prefix_ids();
444
445 tree.insert(id0, commitment0).unwrap();
446 assert_eq!(tree.get(id0), commitment0);
447
448 let err = tree.insert(id1, commitment1).unwrap_err();
449
450 assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
451 duplicate_prefix
452 } if duplicate_prefix == id0.prefix());
453 }
454
455 #[test]
456 fn insert_succeeds_on_multiple_updates() {
457 let mut tree = AccountTree::<Smt>::default();
458 let [(id0, commitment0), (_, commitment1)] = setup_duplicate_prefix_ids();
459
460 tree.insert(id0, commitment0).unwrap();
461 tree.insert(id0, commitment1).unwrap();
462 assert_eq!(tree.get(id0), commitment1);
463 }
464
465 #[test]
466 fn apply_mutations() {
467 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
468 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
469 let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
470
471 let digest0 = Word::from([0, 0, 0, 1u32]);
472 let digest1 = Word::from([0, 0, 0, 2u32]);
473 let digest2 = Word::from([0, 0, 0, 3u32]);
474 let digest3 = Word::from([0, 0, 0, 4u32]);
475
476 let mut tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap();
477
478 let mutations = tree
479 .compute_mutations([(id0, digest1), (id1, digest2), (id2, digest3)])
480 .unwrap();
481
482 tree.apply_mutations(mutations).unwrap();
483
484 assert_eq!(tree.num_accounts(), 3);
485 assert_eq!(tree.get(id0), digest1);
486 assert_eq!(tree.get(id1), digest2);
487 assert_eq!(tree.get(id2), digest3);
488 }
489
490 #[test]
491 fn duplicates_in_compute_mutations() {
492 let [pair0, pair1] = setup_duplicate_prefix_ids();
493 let id2 = AccountIdBuilder::new().build_with_seed([5; 32]);
494 let commitment2 = Word::from([0, 0, 0, 99u32]);
495
496 let tree = AccountTree::with_entries([pair0, (id2, commitment2)]).unwrap();
497 let err = tree.compute_mutations([pair1]).unwrap_err();
498
499 assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
500 duplicate_prefix
501 } if duplicate_prefix == pair1.0.prefix());
502 }
503
504 #[test]
505 fn account_commitments() {
506 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
507 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
508 let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
509
510 let digest0 = Word::from([0, 0, 0, 1u32]);
511 let digest1 = Word::from([0, 0, 0, 2u32]);
512 let digest2 = Word::from([0, 0, 0, 3u32]);
513 let empty_digest = Word::empty();
514
515 let mut tree =
516 AccountTree::with_entries([(id0, digest0), (id1, digest1), (id2, digest2)]).unwrap();
517
518 tree.insert(id2, empty_digest).unwrap();
520
521 assert_eq!(tree.num_accounts(), 2);
522
523 let accounts: Vec<_> = tree.account_commitments().collect();
524 assert_eq!(accounts.len(), 2);
525 assert!(accounts.contains(&(id0, digest0)));
526 assert!(accounts.contains(&(id1, digest1)));
527 }
528
529 #[test]
530 fn account_witness() {
531 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
532 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
533
534 let digest0 = Word::from([0, 0, 0, 1u32]);
535 let digest1 = Word::from([0, 0, 0, 2u32]);
536
537 let tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap();
538
539 assert_eq!(tree.num_accounts(), 2);
540
541 for id in [id0, id1] {
542 let proof = tree.smt.open(&AccountIdKey::from(id).as_word());
543 let (control_path, control_leaf) = proof.into_parts();
544 let witness = tree.open(id);
545
546 assert_eq!(witness.leaf(), control_leaf);
547 assert_eq!(witness.path(), &control_path);
548 }
549 }
550
551 #[test]
552 fn contains_account_prefix() {
553 let [pair0, pair1] = setup_duplicate_prefix_ids();
555 let tree = AccountTree::with_entries([pair0]).unwrap();
556 assert_eq!(tree.num_accounts(), 1);
557
558 assert!(tree.contains_account_id_prefix(pair0.0.prefix()));
560
561 assert!(tree.contains_account_id_prefix(pair1.0.prefix()));
563
564 let id1 = AccountIdBuilder::new().build_with_seed([7; 32]);
566 assert!(!tree.contains_account_id_prefix(id1.prefix()));
567 }
568
569 #[cfg(feature = "std")]
570 #[test]
571 fn large_smt_backend_basic_operations() {
572 use miden_crypto::merkle::smt::{LargeSmt, MemoryStorage};
573
574 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
576 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
577 let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
578
579 let digest0 = Word::from([0, 0, 0, 1u32]);
580 let digest1 = Word::from([0, 0, 0, 2u32]);
581 let digest2 = Word::from([0, 0, 0, 3u32]);
582
583 let tree = LargeSmt::<MemoryStorage>::with_entries(
585 MemoryStorage::default(),
586 [
587 (AccountIdKey::from(id0).as_word(), digest0),
588 (AccountIdKey::from(id1).as_word(), digest1),
589 ],
590 )
591 .map(AccountTree::new_unchecked)
592 .unwrap();
593
594 assert_eq!(tree.num_accounts(), 2);
596 assert_eq!(tree.get(id0), digest0);
597 assert_eq!(tree.get(id1), digest1);
598
599 let witness0 = tree.open(id0);
601 assert_eq!(witness0.id(), id0);
602
603 let mut tree_mut = LargeSmt::<MemoryStorage>::with_entries(
605 MemoryStorage::default(),
606 [
607 (AccountIdKey::from(id0).as_word(), digest0),
608 (AccountIdKey::from(id1).as_word(), digest1),
609 ],
610 )
611 .map(AccountTree::new_unchecked)
612 .unwrap();
613 tree_mut.insert(id2, digest2).unwrap();
614 assert_eq!(tree_mut.num_accounts(), 3);
615 assert_eq!(tree_mut.get(id2), digest2);
616
617 assert_eq!(tree.num_accounts(), 2);
619 }
620
621 #[cfg(feature = "std")]
622 #[test]
623 fn large_smt_backend_duplicate_prefix_check() {
624 use miden_crypto::merkle::smt::{LargeSmt, MemoryStorage};
625
626 let [(id0, commitment0), (id1, commitment1)] = setup_duplicate_prefix_ids();
627
628 let mut tree = AccountTree::new_unchecked(LargeSmt::new(MemoryStorage::default()).unwrap());
629
630 tree.insert(id0, commitment0).unwrap();
631 assert_eq!(tree.get(id0), commitment0);
632
633 let err = tree.insert(id1, commitment1).unwrap_err();
634
635 assert_matches!(
636 err,
637 AccountTreeError::DuplicateIdPrefix { duplicate_prefix }
638 if duplicate_prefix == id0.prefix()
639 );
640 }
641
642 #[cfg(feature = "std")]
643 #[test]
644 fn large_smt_backend_apply_mutations() {
645 use miden_crypto::merkle::smt::{LargeSmt, MemoryStorage};
646
647 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
648 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
649 let id2 = AccountIdBuilder::new().build_with_seed([7; 32]);
650
651 let digest0 = Word::from([0, 0, 0, 1u32]);
652 let digest1 = Word::from([0, 0, 0, 2u32]);
653 let digest2 = Word::from([0, 0, 0, 3u32]);
654 let digest3 = Word::from([0, 0, 0, 4u32]);
655
656 let mut tree = LargeSmt::with_entries(
657 MemoryStorage::default(),
658 [
659 (AccountIdKey::from(id0).as_word(), digest0),
660 (AccountIdKey::from(id1).as_word(), digest1),
661 ],
662 )
663 .map(AccountTree::new_unchecked)
664 .unwrap();
665
666 let mutations = tree
667 .compute_mutations([(id0, digest1), (id1, digest2), (id2, digest3)])
668 .unwrap();
669
670 tree.apply_mutations(mutations).unwrap();
671
672 assert_eq!(tree.num_accounts(), 3);
673 assert_eq!(tree.get(id0), digest1);
674 assert_eq!(tree.get(id1), digest2);
675 assert_eq!(tree.get(id2), digest3);
676 }
677
678 #[cfg(feature = "std")]
679 #[test]
680 fn large_smt_backend_same_root_as_regular_smt() {
681 use miden_crypto::merkle::smt::{LargeSmt, MemoryStorage};
682
683 let id0 = AccountIdBuilder::new().build_with_seed([5; 32]);
684 let id1 = AccountIdBuilder::new().build_with_seed([6; 32]);
685
686 let digest0 = Word::from([0, 0, 0, 1u32]);
687 let digest1 = Word::from([0, 0, 0, 2u32]);
688
689 let large_tree = LargeSmt::with_entries(
691 MemoryStorage::default(),
692 [
693 (AccountIdKey::from(id0).as_word(), digest0),
694 (AccountIdKey::from(id1).as_word(), digest1),
695 ],
696 )
697 .map(AccountTree::new_unchecked)
698 .unwrap();
699
700 let regular_tree = AccountTree::with_entries([(id0, digest0), (id1, digest1)]).unwrap();
702
703 assert_eq!(large_tree.root(), regular_tree.root());
705
706 let large_commitments: std::collections::BTreeMap<_, _> =
708 large_tree.account_commitments().collect();
709 let regular_commitments: std::collections::BTreeMap<_, _> =
710 regular_tree.account_commitments().collect();
711
712 assert_eq!(large_commitments, regular_commitments);
713 }
714}