Skip to main content

miden_protocol/block/account_tree/
backend.rs

1use alloc::boxed::Box;
2use alloc::vec::Vec;
3
4use super::{AccountId, AccountIdKey, AccountIdPrefix, AccountTree, AccountTreeError};
5use crate::Word;
6use crate::crypto::merkle::MerkleError;
7#[cfg(feature = "std")]
8use crate::crypto::merkle::smt::{LargeSmt, LargeSmtError, SmtStorage};
9use crate::crypto::merkle::smt::{LeafIndex, MutationSet, SMT_DEPTH, Smt, SmtLeaf, SmtProof};
10
11// ACCOUNT TREE BACKEND
12// ================================================================================================
13
14/// This trait abstracts over different SMT backends (e.g., `Smt` and `LargeSmt`) to allow
15/// the `AccountTree` to work with either implementation transparently.
16///
17/// Implementors must provide `Default` for creating empty instances. Users should
18/// instantiate the backend directly (potentially with entries) and then pass it to
19/// [`AccountTree::new`].
20pub trait AccountTreeBackend: Sized {
21    type Error: core::error::Error + Send + 'static;
22
23    /// Returns the number of leaves in the SMT.
24    fn num_leaves(&self) -> usize;
25
26    /// Returns all leaves in the SMT as an iterator over leaf index and leaf pairs.
27    fn leaves<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = (LeafIndex<SMT_DEPTH>, SmtLeaf)>>;
28
29    /// Opens the leaf at the given key, returning a Merkle proof.
30    fn open(&self, key: &Word) -> SmtProof;
31
32    /// Applies the given mutation set to the SMT.
33    fn apply_mutations(
34        &mut self,
35        set: MutationSet<SMT_DEPTH, Word, Word>,
36    ) -> Result<(), Self::Error>;
37
38    /// Applies the given mutation set to the SMT and returns the reverse mutation set.
39    ///
40    /// The reverse mutation set can be used to revert the changes made by this operation.
41    fn apply_mutations_with_reversion(
42        &mut self,
43        set: MutationSet<SMT_DEPTH, Word, Word>,
44    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error>;
45
46    /// Computes the mutation set required to apply the given updates to the SMT.
47    fn compute_mutations(
48        &self,
49        updates: Vec<(Word, Word)>,
50    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error>;
51
52    /// Inserts a key-value pair into the SMT, returning the previous value at that key.
53    fn insert(&mut self, key: Word, value: Word) -> Result<Word, Self::Error>;
54
55    /// Returns the value associated with the given key.
56    fn get_value(&self, key: &Word) -> Word;
57
58    /// Returns the leaf at the given key.
59    fn get_leaf(&self, key: &Word) -> SmtLeaf;
60
61    /// Returns the root of the SMT.
62    fn root(&self) -> Word;
63}
64
65// BACKEND IMPLEMENTATION FOR SMT
66// ================================================================================================
67
68impl AccountTreeBackend for Smt {
69    type Error = MerkleError;
70
71    fn num_leaves(&self) -> usize {
72        Smt::num_leaves(self)
73    }
74
75    fn leaves<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = (LeafIndex<SMT_DEPTH>, SmtLeaf)>> {
76        Box::new(Smt::leaves(self).map(|(idx, leaf)| (idx, leaf.clone())))
77    }
78
79    fn open(&self, key: &Word) -> SmtProof {
80        Smt::open(self, key)
81    }
82
83    fn apply_mutations(
84        &mut self,
85        set: MutationSet<SMT_DEPTH, Word, Word>,
86    ) -> Result<(), Self::Error> {
87        Smt::apply_mutations(self, set)
88    }
89
90    fn apply_mutations_with_reversion(
91        &mut self,
92        set: MutationSet<SMT_DEPTH, Word, Word>,
93    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
94        Smt::apply_mutations_with_reversion(self, set)
95    }
96
97    fn compute_mutations(
98        &self,
99        updates: Vec<(Word, Word)>,
100    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
101        Smt::compute_mutations(self, updates)
102    }
103
104    fn insert(&mut self, key: Word, value: Word) -> Result<Word, Self::Error> {
105        Smt::insert(self, key, value)
106    }
107
108    fn get_value(&self, key: &Word) -> Word {
109        Smt::get_value(self, key)
110    }
111
112    fn get_leaf(&self, key: &Word) -> SmtLeaf {
113        Smt::get_leaf(self, key)
114    }
115
116    fn root(&self) -> Word {
117        Smt::root(self)
118    }
119}
120
121// BACKEND IMPLEMENTATION FOR LARGE SMT
122// ================================================================================================
123
124#[cfg(feature = "std")]
125impl<Backend> AccountTreeBackend for LargeSmt<Backend>
126where
127    Backend: SmtStorage,
128{
129    type Error = MerkleError;
130
131    fn num_leaves(&self) -> usize {
132        LargeSmt::num_leaves(self)
133    }
134
135    fn leaves<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = (LeafIndex<SMT_DEPTH>, SmtLeaf)>> {
136        Box::new(LargeSmt::leaves(self).expect("Only IO can error out here"))
137    }
138
139    fn open(&self, key: &Word) -> SmtProof {
140        LargeSmt::open(self, key)
141    }
142
143    fn apply_mutations(
144        &mut self,
145        set: MutationSet<SMT_DEPTH, Word, Word>,
146    ) -> Result<(), Self::Error> {
147        LargeSmt::apply_mutations(self, set).map_err(large_smt_error_to_merkle_error)
148    }
149
150    fn apply_mutations_with_reversion(
151        &mut self,
152        set: MutationSet<SMT_DEPTH, Word, Word>,
153    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
154        LargeSmt::apply_mutations_with_reversion(self, set).map_err(large_smt_error_to_merkle_error)
155    }
156
157    fn compute_mutations(
158        &self,
159        updates: Vec<(Word, Word)>,
160    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
161        LargeSmt::compute_mutations(self, updates).map_err(large_smt_error_to_merkle_error)
162    }
163
164    fn insert(&mut self, key: Word, value: Word) -> Result<Word, Self::Error> {
165        LargeSmt::insert(self, key, value)
166    }
167
168    fn get_value(&self, key: &Word) -> Word {
169        LargeSmt::get_value(self, key)
170    }
171
172    fn get_leaf(&self, key: &Word) -> SmtLeaf {
173        LargeSmt::get_leaf(self, key)
174    }
175
176    fn root(&self) -> Word {
177        LargeSmt::root(self)
178    }
179}
180
181// CONVENIENCE METHODS
182// ================================================================================================
183
184impl AccountTree<Smt> {
185    /// Creates a new [`AccountTree`] with the provided entries.
186    ///
187    /// This is a convenience method for testing that creates an SMT backend with the provided
188    /// entries and wraps it in an AccountTree. It validates that the entries don't contain
189    /// duplicate prefixes.
190    ///
191    /// # Errors
192    ///
193    /// Returns an error if:
194    /// - The provided entries contain duplicate account ID prefixes
195    /// - The backend fails to create the SMT with the entries
196    pub fn with_entries<I>(
197        entries: impl IntoIterator<Item = (AccountId, Word), IntoIter = I>,
198    ) -> Result<Self, AccountTreeError>
199    where
200        I: ExactSizeIterator<Item = (AccountId, Word)>,
201    {
202        // Create the SMT with the entries
203        let smt = Smt::with_entries(
204            entries
205                .into_iter()
206                .map(|(id, commitment)| (AccountIdKey::from(id).as_word(), commitment)),
207        )
208        .map_err(|err| {
209            let MerkleError::DuplicateValuesForIndex(leaf_idx) = err else {
210                unreachable!("the only error returned by Smt::with_entries is of this type");
211            };
212
213            // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed that
214            // the leaf_idx is a valid Felt as well as a valid account ID prefix.
215            AccountTreeError::DuplicateStateCommitments {
216                prefix: AccountIdPrefix::new_unchecked(
217                    crate::Felt::try_from(leaf_idx).expect("leaf index should be a valid felt"),
218                ),
219            }
220        })?;
221
222        AccountTree::new(smt)
223    }
224}
225
226// HELPER FUNCTIONS
227// ================================================================================================
228
229#[cfg(feature = "std")]
230fn large_smt_error_to_merkle_error(err: LargeSmtError) -> MerkleError {
231    match err {
232        LargeSmtError::Storage(storage_err) => {
233            panic!("Storage error encountered: {:?}", storage_err)
234        },
235        LargeSmtError::StorageNotEmpty => {
236            panic!("StorageNotEmpty error encountered: {:?}", err)
237        },
238        LargeSmtError::Merkle(merkle_err) => merkle_err,
239        LargeSmtError::RootMismatch { expected, actual } => MerkleError::ConflictingRoots {
240            expected_root: expected,
241            actual_root: actual,
242        },
243    }
244}