miden-crypto/src/merkle/smt/full/mod.rs
src/merkle/smt/full/mod.rsRaw
1use alloc::{
2 collections::{BTreeMap, BTreeSet},
3 string::ToString,
4 vec::Vec,
5};
6
7use super::{
8 EmptySubtreeRoots, Felt, InnerNode, InnerNodeInfo, LeafIndex, MerkleError, MerklePath,
9 MutationSet, NodeIndex, Rpo256, RpoDigest, SparseMerkleTree, Word, EMPTY_WORD,
10};
11
12mod error;
13pub use error::{SmtLeafError, SmtProofError};
14
15mod leaf;
16pub use leaf::SmtLeaf;
17
18mod proof;
19pub use proof::SmtProof;
20use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
21
22#[cfg(test)]
23mod tests;
24
25// CONSTANTS
26// ================================================================================================
27
28pub const SMT_DEPTH: u8 = 64;
29
30// SMT
31// ================================================================================================
32
33/// Sparse Merkle tree mapping 256-bit keys to 256-bit values. Both keys and values are represented
34/// by 4 field elements.
35///
36/// All leaves sit at depth 64. The most significant element of the key is used to identify the leaf
37/// to which the key maps.
38///
39/// A leaf is either empty, or holds one or more key-value pairs. An empty leaf hashes to the empty
40/// word. Otherwise, a leaf hashes to the hash of its key-value pairs, ordered by key first, value
41/// second.
42#[derive(Debug, Clone, PartialEq, Eq)]
43#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
44pub struct Smt {
45 root: RpoDigest,
46 leaves: BTreeMap<u64, SmtLeaf>,
47 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
48}
49
50impl Smt {
51 // CONSTANTS
52 // --------------------------------------------------------------------------------------------
53 /// The default value used to compute the hash of empty leaves
54 pub const EMPTY_VALUE: Word = <Self as SparseMerkleTree<SMT_DEPTH>>::EMPTY_VALUE;
55
56 // CONSTRUCTORS
57 // --------------------------------------------------------------------------------------------
58
59 /// Returns a new [Smt].
60 ///
61 /// All leaves in the returned tree are set to [Self::EMPTY_VALUE].
62 pub fn new() -> Self {
63 let root = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
64
65 Self {
66 root,
67 leaves: BTreeMap::new(),
68 inner_nodes: BTreeMap::new(),
69 }
70 }
71
72 /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
73 ///
74 /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
75 ///
76 /// # Errors
77 /// Returns an error if the provided entries contain multiple values for the same key.
78 pub fn with_entries(
79 entries: impl IntoIterator<Item = (RpoDigest, Word)>,
80 ) -> Result<Self, MerkleError> {
81 // create an empty tree
82 let mut tree = Self::new();
83
84 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
85 // entries with the empty value need additional tracking.
86 let mut key_set_to_zero = BTreeSet::new();
87
88 for (key, value) in entries {
89 let old_value = tree.insert(key, value);
90
91 if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
92 return Err(MerkleError::DuplicateValuesForIndex(
93 LeafIndex::<SMT_DEPTH>::from(key).value(),
94 ));
95 }
96
97 if value == EMPTY_WORD {
98 key_set_to_zero.insert(key);
99 };
100 }
101 Ok(tree)
102 }
103
104 // PUBLIC ACCESSORS
105 // --------------------------------------------------------------------------------------------
106
107 /// Returns the depth of the tree
108 pub const fn depth(&self) -> u8 {
109 SMT_DEPTH
110 }
111
112 /// Returns the root of the tree
113 pub fn root(&self) -> RpoDigest {
114 <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
115 }
116
117 /// Returns the number of non-empty leaves in this tree.
118 pub fn num_leaves(&self) -> usize {
119 self.leaves.len()
120 }
121
122 /// Returns the leaf to which `key` maps
123 pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
124 <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
125 }
126
127 /// Returns the value associated with `key`
128 pub fn get_value(&self, key: &RpoDigest) -> Word {
129 <Self as SparseMerkleTree<SMT_DEPTH>>::get_value(self, key)
130 }
131
132 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
133 /// path to the leaf, as well as the leaf itself.
134 pub fn open(&self, key: &RpoDigest) -> SmtProof {
135 <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
136 }
137
138 /// Returns a boolean value indicating whether the SMT is empty.
139 pub fn is_empty(&self) -> bool {
140 debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);
141 self.root == Self::EMPTY_ROOT
142 }
143
144 // ITERATORS
145 // --------------------------------------------------------------------------------------------
146
147 /// Returns an iterator over the leaves of this [Smt].
148 pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
149 self.leaves
150 .iter()
151 .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
152 }
153
154 /// Returns an iterator over the key-value pairs of this [Smt].
155 pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
156 self.leaves().flat_map(|(_, leaf)| leaf.entries())
157 }
158
159 /// Returns an iterator over the inner nodes of this [Smt].
160 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
161 self.inner_nodes.values().map(|e| InnerNodeInfo {
162 value: e.hash(),
163 left: e.left,
164 right: e.right,
165 })
166 }
167
168 // STATE MUTATORS
169 // --------------------------------------------------------------------------------------------
170
171 /// Inserts a value at the specified key, returning the previous value associated with that key.
172 /// Recall that by definition, any key that hasn't been updated is associated with
173 /// [`Self::EMPTY_VALUE`].
174 ///
175 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
176 /// updating the root itself.
177 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
178 <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
179 }
180
181 /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
182 /// tree, allowing for validation before applying those changes.
183 ///
184 /// This method returns a [`MutationSet`], which contains all the information for inserting
185 /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
186 /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
187 /// [`Smt::apply_mutations()`] can be called in order to commit these changes to the Merkle
188 /// tree, or [`drop()`] to discard them.
189 ///
190 /// # Example
191 /// ```
192 /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
193 /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
194 /// let mut smt = Smt::new();
195 /// let pair = (RpoDigest::default(), Word::default());
196 /// let mutations = smt.compute_mutations(vec![pair]);
197 /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
198 /// smt.apply_mutations(mutations);
199 /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
200 /// ```
201 pub fn compute_mutations(
202 &self,
203 kv_pairs: impl IntoIterator<Item = (RpoDigest, Word)>,
204 ) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
205 <Self as SparseMerkleTree<SMT_DEPTH>>::compute_mutations(self, kv_pairs)
206 }
207
208 /// Applies the prospective mutations computed with [`Smt::compute_mutations()`] to this tree.
209 ///
210 /// # Errors
211 /// If `mutations` was computed on a tree with a different root than this one, returns
212 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
213 /// the `mutations` were computed against, and the second item is the actual current root of
214 /// this tree.
215 pub fn apply_mutations(
216 &mut self,
217 mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
218 ) -> Result<(), MerkleError> {
219 <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations(self, mutations)
220 }
221
222 /// Applies the prospective mutations computed with [`Smt::compute_mutations()`] to this tree
223 /// and returns the reverse mutation set.
224 ///
225 /// Applying the reverse mutation sets to the updated tree will revert the changes.
226 ///
227 /// # Errors
228 /// If `mutations` was computed on a tree with a different root than this one, returns
229 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
230 /// the `mutations` were computed against, and the second item is the actual current root of
231 /// this tree.
232 pub fn apply_mutations_with_reversion(
233 &mut self,
234 mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
235 ) -> Result<MutationSet<SMT_DEPTH, RpoDigest, Word>, MerkleError> {
236 <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations_with_reversion(self, mutations)
237 }
238
239 // HELPERS
240 // --------------------------------------------------------------------------------------------
241
242 /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
243 /// value, such that this is indeed an insertion.
244 fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
245 debug_assert_ne!(value, Self::EMPTY_VALUE);
246
247 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
248
249 match self.leaves.get_mut(&leaf_index.value()) {
250 Some(leaf) => leaf.insert(key, value),
251 None => {
252 self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
253
254 None
255 },
256 }
257 }
258
259 /// Removes key-value pair at leaf index pointed to by `key` if it exists.
260 fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
261 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
262
263 if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
264 let (old_value, is_empty) = leaf.remove(key);
265 if is_empty {
266 self.leaves.remove(&leaf_index.value());
267 }
268 old_value
269 } else {
270 // there's nothing stored at the leaf; nothing to update
271 None
272 }
273 }
274}
275
276impl SparseMerkleTree<SMT_DEPTH> for Smt {
277 type Key = RpoDigest;
278 type Value = Word;
279 type Leaf = SmtLeaf;
280 type Opening = SmtProof;
281
282 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
283 const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
284
285 fn root(&self) -> RpoDigest {
286 self.root
287 }
288
289 fn set_root(&mut self, root: RpoDigest) {
290 self.root = root;
291 }
292
293 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
294 self.inner_nodes
295 .get(&index)
296 .cloned()
297 .unwrap_or_else(|| EmptySubtreeRoots::get_inner_node(SMT_DEPTH, index.depth()))
298 }
299
300 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) -> Option<InnerNode> {
301 self.inner_nodes.insert(index, inner_node)
302 }
303
304 fn remove_inner_node(&mut self, index: NodeIndex) -> Option<InnerNode> {
305 self.inner_nodes.remove(&index)
306 }
307
308 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
309 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
310 if value != Self::EMPTY_VALUE {
311 self.perform_insert(key, value)
312 } else {
313 self.perform_remove(key)
314 }
315 }
316
317 fn get_value(&self, key: &Self::Key) -> Self::Value {
318 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
319
320 match self.leaves.get(&leaf_pos) {
321 Some(leaf) => leaf.get_value(key).unwrap_or_default(),
322 None => EMPTY_WORD,
323 }
324 }
325
326 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
327 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
328
329 match self.leaves.get(&leaf_pos) {
330 Some(leaf) => leaf.clone(),
331 None => SmtLeaf::new_empty(key.into()),
332 }
333 }
334
335 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
336 leaf.hash()
337 }
338
339 fn construct_prospective_leaf(
340 &self,
341 mut existing_leaf: SmtLeaf,
342 key: &RpoDigest,
343 value: &Word,
344 ) -> SmtLeaf {
345 debug_assert_eq!(existing_leaf.index(), Self::key_to_leaf_index(key));
346
347 match existing_leaf {
348 SmtLeaf::Empty(_) => SmtLeaf::new_single(*key, *value),
349 _ => {
350 if *value != EMPTY_WORD {
351 existing_leaf.insert(*key, *value);
352 } else {
353 existing_leaf.remove(*key);
354 }
355
356 existing_leaf
357 },
358 }
359 }
360
361 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
362 let most_significant_felt = key[3];
363 LeafIndex::new_max_depth(most_significant_felt.as_int())
364 }
365
366 fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
367 SmtProof::new_unchecked(path, leaf)
368 }
369}
370
371impl Default for Smt {
372 fn default() -> Self {
373 Self::new()
374 }
375}
376
377// CONVERSIONS
378// ================================================================================================
379
380impl From<Word> for LeafIndex<SMT_DEPTH> {
381 fn from(value: Word) -> Self {
382 // We use the most significant `Felt` of a `Word` as the leaf index.
383 Self::new_max_depth(value[3].as_int())
384 }
385}
386
387impl From<RpoDigest> for LeafIndex<SMT_DEPTH> {
388 fn from(value: RpoDigest) -> Self {
389 Word::from(value).into()
390 }
391}
392
393impl From<&RpoDigest> for LeafIndex<SMT_DEPTH> {
394 fn from(value: &RpoDigest) -> Self {
395 Word::from(value).into()
396 }
397}
398
399// SERIALIZATION
400// ================================================================================================
401
402impl Serializable for Smt {
403 fn write_into<W: ByteWriter>(&self, target: &mut W) {
404 // Write the number of filled leaves for this Smt
405 target.write_usize(self.entries().count());
406
407 // Write each (key, value) pair
408 for (key, value) in self.entries() {
409 target.write(key);
410 target.write(value);
411 }
412 }
413
414 fn get_size_hint(&self) -> usize {
415 let entries_count = self.entries().count();
416
417 // Each entry is the size of a digest plus a word.
418 entries_count.get_size_hint()
419 + entries_count * (RpoDigest::SERIALIZED_SIZE + EMPTY_WORD.get_size_hint())
420 }
421}
422
423impl Deserializable for Smt {
424 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
425 // Read the number of filled leaves for this Smt
426 let num_filled_leaves = source.read_usize()?;
427 let mut entries = Vec::with_capacity(num_filled_leaves);
428
429 for _ in 0..num_filled_leaves {
430 let key = source.read()?;
431 let value = source.read()?;
432 entries.push((key, value));
433 }
434
435 Self::with_entries(entries)
436 .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
437 }
438}
439
440#[test]
441fn test_smt_serialization_deserialization() {
442 // Smt for default types (empty map)
443 let smt_default = Smt::default();
444 let bytes = smt_default.to_bytes();
445 assert_eq!(smt_default, Smt::read_from_bytes(&bytes).unwrap());
446 assert_eq!(bytes.len(), smt_default.get_size_hint());
447
448 // Smt with values
449 let smt_leaves_2: [(RpoDigest, Word); 2] = [
450 (
451 RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
452 [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
453 ),
454 (
455 RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]),
456 [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
457 ),
458 ];
459 let smt = Smt::with_entries(smt_leaves_2).unwrap();
460
461 let bytes = smt.to_bytes();
462 assert_eq!(smt, Smt::read_from_bytes(&bytes).unwrap());
463 assert_eq!(bytes.len(), smt.get_size_hint());
464}