miden_objects/block/
partial_account_tree.rs

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

miden_objects/block/
partial_account_tree.rs
