miden_objects/block/
partial_account_tree.rs

1use miden_crypto::merkle::SmtLeaf;
2
3use crate::Word;
4use crate::account::AccountId;
5use crate::block::AccountWitness;
6use crate::block::account_tree::account_id_to_smt_key;
7use crate::crypto::merkle::PartialSmt;
8use crate::errors::AccountTreeError;
9
10/// The partial sparse merkle tree containing the state commitments of accounts in the chain.
11///
12/// This is the partial version of [`AccountTree`](crate::block::account_tree::AccountTree).
13#[derive(Debug, Clone, Default, PartialEq, Eq)]
14pub struct PartialAccountTree {
15    smt: PartialSmt,
16}
17
18impl PartialAccountTree {
19    // CONSTRUCTORS
20    // --------------------------------------------------------------------------------------------
21
22    /// Creates a new partial account tree with the provided root that does not track any account
23    /// IDs.
24    pub fn new(root: Word) -> Self {
25        PartialAccountTree { smt: PartialSmt::new(root) }
26    }
27
28    /// Returns a new [`PartialAccountTree`] instantiated with the provided entries.
29    ///
30    /// # Errors
31    ///
32    /// Returns an error if:
33    /// - the merkle paths of the witnesses do not result in the same tree root.
34    /// - there are multiple witnesses for the same ID _prefix_.
35    pub fn with_witnesses(
36        witnesses: impl IntoIterator<Item = AccountWitness>,
37    ) -> Result<Self, AccountTreeError> {
38        let mut witnesses = witnesses.into_iter();
39
40        let Some(first_witness) = witnesses.next() else {
41            return Ok(Self::default());
42        };
43
44        // Construct a partial account tree with the root of the first witness.
45        // SAFETY: This is guaranteed to _not_ result in a tree with more than one entry because
46        // the account witness type guarantees that it tracks zero or one entries.
47        let partial_smt = PartialSmt::from_proofs([first_witness.into_proof()])
48            .map_err(AccountTreeError::TreeRootConflict)?;
49        let mut tree = PartialAccountTree { smt: partial_smt };
50
51        // Add all remaining witnesses to the tree, which validates the invariants of the account
52        // tree.
53        for witness in witnesses {
54            tree.track_account(witness)?;
55        }
56
57        Ok(tree)
58    }
59
60    // PUBLIC ACCESSORS
61    // --------------------------------------------------------------------------------------------
62
63    /// Returns an opening of the leaf associated with the `account_id`. This is a proof of the
64    /// current state commitment of the given account ID.
65    ///
66    /// Conceptually, an opening is a Merkle path to the leaf, as well as the leaf itself.
67    ///
68    /// # Errors
69    ///
70    /// Returns an error if:
71    /// - the account ID is not tracked by this account tree.
72    pub fn open(&self, account_id: AccountId) -> Result<AccountWitness, AccountTreeError> {
73        let key = account_id_to_smt_key(account_id);
74
75        self.smt
76            .open(&key)
77            .map(|proof| AccountWitness::from_smt_proof(account_id, proof))
78            .map_err(|source| AccountTreeError::UntrackedAccountId { id: account_id, source })
79    }
80
81    /// Returns the current state commitment of the given account ID.
82    ///
83    /// # Errors
84    ///
85    /// Returns an error if:
86    /// - the account ID is not tracked by this account tree.
87    pub fn get(&self, account_id: AccountId) -> Result<Word, AccountTreeError> {
88        let key = account_id_to_smt_key(account_id);
89        self.smt
90            .get_value(&key)
91            .map_err(|source| AccountTreeError::UntrackedAccountId { id: account_id, source })
92    }
93
94    /// Returns the root of the tree.
95    pub fn root(&self) -> Word {
96        self.smt.root()
97    }
98
99    // PUBLIC MUTATORS
100    // --------------------------------------------------------------------------------------------
101
102    /// Adds the given account witness to the partial tree and tracks it. Once an account has
103    /// been added to the tree, it can be updated using [`Self::upsert_state_commitments`].
104    ///
105    /// # Errors
106    ///
107    /// Returns an error if:
108    /// - after the witness' merkle path was added, the partial account tree has a different root
109    ///   than before it was added (except when the first witness is added).
110    /// - there exists a leaf in the tree whose account ID prefix matches the one in the provided
111    ///   witness.
112    pub fn track_account(&mut self, witness: AccountWitness) -> Result<(), AccountTreeError> {
113        let id_prefix = witness.id().prefix();
114        let id_key = account_id_to_smt_key(witness.id());
115
116        // If a leaf with the same prefix is already tracked by this partial tree, consider it an
117        // error.
118        //
119        // We return an error even for empty leaves, because tracking the same ID prefix twice
120        // indicates that different IDs are attempted to be tracked. It would technically not
121        // violate the invariant of the tree that it only tracks zero or one entries per leaf, but
122        // since tracking the same ID twice should practically never happen, we return an error, out
123        // of an abundance of caution.
124        if self.smt.get_leaf(&id_key).is_ok() {
125            return Err(AccountTreeError::DuplicateIdPrefix { duplicate_prefix: id_prefix });
126        }
127
128        self.smt
129            .add_proof(witness.into_proof())
130            .map_err(AccountTreeError::TreeRootConflict)?;
131
132        Ok(())
133    }
134
135    /// Inserts or updates the provided account ID -> state commitment updates into the partial tree
136    /// which results in a new tree root.
137    ///
138    /// # Errors
139    ///
140    /// Returns an error if:
141    /// - the prefix of the account ID already exists in the tree.
142    /// - the account_id is not tracked by this partial account tree.
143    pub fn upsert_state_commitments(
144        &mut self,
145        updates: impl IntoIterator<Item = (AccountId, Word)>,
146    ) -> Result<(), AccountTreeError> {
147        for (account_id, state_commitment) in updates {
148            self.insert(account_id, state_commitment)?;
149        }
150
151        Ok(())
152    }
153
154    /// Inserts the state commitment for the given account ID, returning the previous state
155    /// commitment associated with that ID.
156    ///
157    /// This also recomputes all hashes between the leaf (associated with the key) and the root,
158    /// updating the root itself.
159    ///
160    /// # Errors
161    ///
162    /// Returns an error if:
163    /// - the prefix of the account ID already exists in the tree.
164    /// - the account_id is not tracked by this partial account tree.
165    fn insert(
166        &mut self,
167        account_id: AccountId,
168        state_commitment: Word,
169    ) -> Result<Word, AccountTreeError> {
170        let key = account_id_to_smt_key(account_id);
171
172        // If there exists a tracked leaf whose key is _not_ the one we're about to overwrite, then
173        // we would insert the new commitment next to an existing account ID with the same prefix,
174        // which is an error.
175        // Note that if the leaf is empty, that's fine. It means it is tracked by the partial SMT,
176        // but no account ID is inserted yet.
177        // Also note that the multiple variant cannot occur by construction of the tree.
178        if let Ok(SmtLeaf::Single((existing_key, _))) = self.smt.get_leaf(&key)
179            && key != existing_key
180        {
181            return Err(AccountTreeError::DuplicateIdPrefix {
182                duplicate_prefix: account_id.prefix(),
183            });
184        }
185
186        self.smt
187            .insert(key, state_commitment)
188            .map_err(|source| AccountTreeError::UntrackedAccountId { id: account_id, source })
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use assert_matches::assert_matches;
195    use miden_crypto::merkle::Smt;
196
197    use super::*;
198    use crate::block::account_tree::AccountTree;
199    use crate::block::account_tree::tests::setup_duplicate_prefix_ids;
200
201    #[test]
202    fn insert_fails_on_duplicate_prefix() -> anyhow::Result<()> {
203        let mut full_tree = AccountTree::<Smt>::default();
204
205        let [(id0, commitment0), (id1, commitment1)] = setup_duplicate_prefix_ids();
206
207        full_tree.insert(id0, commitment0).unwrap();
208        let witness = full_tree.open(id0);
209
210        let mut partial_tree = PartialAccountTree::with_witnesses([witness])?;
211
212        partial_tree.insert(id0, commitment0).unwrap();
213        assert_eq!(partial_tree.get(id0).unwrap(), commitment0);
214
215        let err = partial_tree.insert(id1, commitment1).unwrap_err();
216
217        assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
218          duplicate_prefix
219        } if duplicate_prefix == id0.prefix());
220
221        partial_tree.upsert_state_commitments([(id1, commitment1)]).unwrap_err();
222
223        assert_matches!(err, AccountTreeError::DuplicateIdPrefix {
224          duplicate_prefix
225        } if duplicate_prefix == id0.prefix());
226
227        Ok(())
228    }
229
230    #[test]
231    fn insert_succeeds_on_multiple_updates() {
232        let mut full_tree = AccountTree::<Smt>::default();
233        let [(id0, commitment0), (_, commitment1)] = setup_duplicate_prefix_ids();
234
235        full_tree.insert(id0, commitment0).unwrap();
236        let witness = full_tree.open(id0);
237
238        let mut partial_tree = PartialAccountTree::new(full_tree.root());
239
240        partial_tree.track_account(witness.clone()).unwrap();
241        assert_eq!(
242            partial_tree.open(id0).unwrap(),
243            witness,
244            "full tree witness and partial tree witness should be the same"
245        );
246        assert_eq!(
247            partial_tree.root(),
248            full_tree.root(),
249            "full tree root and partial tree root should be the same"
250        );
251
252        partial_tree.insert(id0, commitment0).unwrap();
253        partial_tree.insert(id0, commitment1).unwrap();
254        assert_eq!(partial_tree.get(id0).unwrap(), commitment1);
255    }
256
257    #[test]
258    fn upsert_state_commitments_fails_on_untracked_key() {
259        let mut partial_tree = PartialAccountTree::default();
260        let [update, _] = setup_duplicate_prefix_ids();
261
262        let err = partial_tree.upsert_state_commitments([update]).unwrap_err();
263        assert_matches!(err, AccountTreeError::UntrackedAccountId { id, .. }
264          if id == update.0
265        )
266    }
267
268    #[test]
269    fn track_fails_on_duplicate_prefix() {
270        // Use a raw Smt since an account tree would not allow us to get the witnesses for two
271        // account IDs with the same prefix.
272        let full_tree = Smt::with_entries(
273            setup_duplicate_prefix_ids()
274                .map(|(id, commitment)| (account_id_to_smt_key(id), commitment)),
275        )
276        .unwrap();
277
278        let [(id0, _), (id1, _)] = setup_duplicate_prefix_ids();
279
280        let key0 = account_id_to_smt_key(id0);
281        let key1 = account_id_to_smt_key(id1);
282        let proof0 = full_tree.open(&key0);
283        let proof1 = full_tree.open(&key1);
284        assert_eq!(proof0.leaf(), proof1.leaf());
285
286        let witness0 =
287            AccountWitness::new(id0, proof0.get(&key0).unwrap(), proof0.into_parts().0).unwrap();
288        let witness1 =
289            AccountWitness::new(id1, proof1.get(&key1).unwrap(), proof1.into_parts().0).unwrap();
290
291        let mut partial_tree = PartialAccountTree::with_witnesses([witness0]).unwrap();
292        let err = partial_tree.track_account(witness1).unwrap_err();
293
294        assert_matches!(err, AccountTreeError::DuplicateIdPrefix { duplicate_prefix, .. }
295          if duplicate_prefix == id1.prefix()
296        )
297    }
298}
miden_objects/block/partial_account_tree.rs

miden_objects/block/
partial_account_tree.rs
