Skip to main content

miden_protocol/block/account_tree/
backend.rs

1use alloc::boxed::Box;
2use alloc::vec::Vec;
3
4use super::{AccountId, AccountIdPrefix, AccountTree, AccountTreeError, account_id_to_smt_key};
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 returns Result<usize, LargeSmtError>
133        // We'll unwrap or return 0 on error
134        LargeSmt::num_leaves(self).map_err(large_smt_error_to_merkle_error).unwrap_or(0)
135    }
136
137    fn leaves<'a>(&'a self) -> Box<dyn 'a + Iterator<Item = (LeafIndex<SMT_DEPTH>, SmtLeaf)>> {
138        Box::new(LargeSmt::leaves(self).expect("Only IO can error out here"))
139    }
140
141    fn open(&self, key: &Word) -> SmtProof {
142        LargeSmt::open(self, key)
143    }
144
145    fn apply_mutations(
146        &mut self,
147        set: MutationSet<SMT_DEPTH, Word, Word>,
148    ) -> Result<(), Self::Error> {
149        LargeSmt::apply_mutations(self, set).map_err(large_smt_error_to_merkle_error)
150    }
151
152    fn apply_mutations_with_reversion(
153        &mut self,
154        set: MutationSet<SMT_DEPTH, Word, Word>,
155    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
156        LargeSmt::apply_mutations_with_reversion(self, set).map_err(large_smt_error_to_merkle_error)
157    }
158
159    fn compute_mutations(
160        &self,
161        updates: Vec<(Word, Word)>,
162    ) -> Result<MutationSet<SMT_DEPTH, Word, Word>, Self::Error> {
163        LargeSmt::compute_mutations(self, updates).map_err(large_smt_error_to_merkle_error)
164    }
165
166    fn insert(&mut self, key: Word, value: Word) -> Result<Word, Self::Error> {
167        LargeSmt::insert(self, key, value)
168    }
169
170    fn get_value(&self, key: &Word) -> Word {
171        LargeSmt::get_value(self, key)
172    }
173
174    fn get_leaf(&self, key: &Word) -> SmtLeaf {
175        LargeSmt::get_leaf(self, key)
176    }
177
178    fn root(&self) -> Word {
179        LargeSmt::root(self)
180    }
181}
182
183// CONVENIENCE METHODS
184// ================================================================================================
185
186impl AccountTree<Smt> {
187    /// Creates a new [`AccountTree`] with the provided entries.
188    ///
189    /// This is a convenience method for testing that creates an SMT backend with the provided
190    /// entries and wraps it in an AccountTree. It validates that the entries don't contain
191    /// duplicate prefixes.
192    ///
193    /// # Errors
194    ///
195    /// Returns an error if:
196    /// - The provided entries contain duplicate account ID prefixes
197    /// - The backend fails to create the SMT with the entries
198    pub fn with_entries<I>(
199        entries: impl IntoIterator<Item = (AccountId, Word), IntoIter = I>,
200    ) -> Result<Self, AccountTreeError>
201    where
202        I: ExactSizeIterator<Item = (AccountId, Word)>,
203    {
204        // Create the SMT with the entries
205        let smt = Smt::with_entries(
206            entries
207                .into_iter()
208                .map(|(id, commitment)| (account_id_to_smt_key(id), commitment)),
209        )
210        .map_err(|err| {
211            let MerkleError::DuplicateValuesForIndex(leaf_idx) = err else {
212                unreachable!("the only error returned by Smt::with_entries is of this type");
213            };
214
215            // SAFETY: Since we only inserted account IDs into the SMT, it is guaranteed that
216            // the leaf_idx is a valid Felt as well as a valid account ID prefix.
217            AccountTreeError::DuplicateStateCommitments {
218                prefix: AccountIdPrefix::new_unchecked(
219                    crate::Felt::try_from(leaf_idx).expect("leaf index should be a valid felt"),
220                ),
221            }
222        })?;
223
224        AccountTree::new(smt)
225    }
226}
227
228// HELPER FUNCTIONS
229// ================================================================================================
230
231#[cfg(feature = "std")]
232fn large_smt_error_to_merkle_error(err: LargeSmtError) -> MerkleError {
233    match err {
234        LargeSmtError::Storage(storage_err) => {
235            panic!("Storage error encountered: {:?}", storage_err)
236        },
237        LargeSmtError::Merkle(merkle_err) => merkle_err,
238    }
239}