miden_objects/block/
partial_account_tree.rs

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

miden_objects/block/
partial_account_tree.rs
