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 /// If the `concurrent` feature is enabled, this function uses a parallel implementation to
75 /// process the entries efficiently, otherwise it defaults to the sequential implementation.
76 ///
77 /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
78 ///
79 /// # Errors
80 /// Returns an error if the provided entries contain multiple values for the same key.
81 pub fn with_entries(
82 entries: impl IntoIterator<Item = (RpoDigest, Word)>,
83 ) -> Result<Self, MerkleError> {
84 #[cfg(feature = "concurrent")]
85 {
86 let mut seen_keys = BTreeSet::new();
87 let entries: Vec<_> = entries
88 .into_iter()
89 .map(|(key, value)| {
90 if seen_keys.insert(key) {
91 Ok((key, value))
92 } else {
93 Err(MerkleError::DuplicateValuesForIndex(
94 LeafIndex::<SMT_DEPTH>::from(key).value(),
95 ))
96 }
97 })
98 .collect::<Result<_, _>>()?;
99 if entries.is_empty() {
100 return Ok(Self::default());
101 }
102 <Self as SparseMerkleTree<SMT_DEPTH>>::with_entries_par(entries)
103 }
104 #[cfg(not(feature = "concurrent"))]
105 {
106 Self::with_entries_sequential(entries)
107 }
108 }
109
110 /// Returns a new [Smt] instantiated with leaves set as specified by the provided entries.
111 ///
112 /// This sequential implementation processes entries one at a time to build the tree.
113 /// All leaves omitted from the entries list are set to [Self::EMPTY_VALUE].
114 ///
115 /// # Errors
116 /// Returns an error if the provided entries contain multiple values for the same key.
117 pub fn with_entries_sequential(
118 entries: impl IntoIterator<Item = (RpoDigest, Word)>,
119 ) -> Result<Self, MerkleError> {
120 // create an empty tree
121 let mut tree = Self::new();
122
123 // This being a sparse data structure, the EMPTY_WORD is not assigned to the `BTreeMap`, so
124 // entries with the empty value need additional tracking.
125 let mut key_set_to_zero = BTreeSet::new();
126
127 for (key, value) in entries {
128 let old_value = tree.insert(key, value);
129
130 if old_value != EMPTY_WORD || key_set_to_zero.contains(&key) {
131 return Err(MerkleError::DuplicateValuesForIndex(
132 LeafIndex::<SMT_DEPTH>::from(key).value(),
133 ));
134 }
135
136 if value == EMPTY_WORD {
137 key_set_to_zero.insert(key);
138 };
139 }
140 Ok(tree)
141 }
142
143 /// Returns a new [`Smt`] instantiated from already computed leaves and nodes.
144 ///
145 /// This function performs minimal consistency checking. It is the caller's responsibility to
146 /// ensure the passed arguments are correct and consistent with each other.
147 ///
148 /// # Panics
149 /// With debug assertions on, this function panics if `root` does not match the root node in
150 /// `inner_nodes`.
151 pub fn from_raw_parts(
152 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
153 leaves: BTreeMap<u64, SmtLeaf>,
154 root: RpoDigest,
155 ) -> Self {
156 // Our particular implementation of `from_raw_parts()` never returns `Err`.
157 <Self as SparseMerkleTree<SMT_DEPTH>>::from_raw_parts(inner_nodes, leaves, root).unwrap()
158 }
159
160 // PUBLIC ACCESSORS
161 // --------------------------------------------------------------------------------------------
162
163 /// Returns the depth of the tree
164 pub const fn depth(&self) -> u8 {
165 SMT_DEPTH
166 }
167
168 /// Returns the root of the tree
169 pub fn root(&self) -> RpoDigest {
170 <Self as SparseMerkleTree<SMT_DEPTH>>::root(self)
171 }
172
173 /// Returns the number of non-empty leaves in this tree.
174 pub fn num_leaves(&self) -> usize {
175 self.leaves.len()
176 }
177
178 /// Returns the leaf to which `key` maps
179 pub fn get_leaf(&self, key: &RpoDigest) -> SmtLeaf {
180 <Self as SparseMerkleTree<SMT_DEPTH>>::get_leaf(self, key)
181 }
182
183 /// Returns the value associated with `key`
184 pub fn get_value(&self, key: &RpoDigest) -> Word {
185 <Self as SparseMerkleTree<SMT_DEPTH>>::get_value(self, key)
186 }
187
188 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
189 /// path to the leaf, as well as the leaf itself.
190 pub fn open(&self, key: &RpoDigest) -> SmtProof {
191 <Self as SparseMerkleTree<SMT_DEPTH>>::open(self, key)
192 }
193
194 /// Returns a boolean value indicating whether the SMT is empty.
195 pub fn is_empty(&self) -> bool {
196 debug_assert_eq!(self.leaves.is_empty(), self.root == Self::EMPTY_ROOT);
197 self.root == Self::EMPTY_ROOT
198 }
199
200 // ITERATORS
201 // --------------------------------------------------------------------------------------------
202
203 /// Returns an iterator over the leaves of this [Smt].
204 pub fn leaves(&self) -> impl Iterator<Item = (LeafIndex<SMT_DEPTH>, &SmtLeaf)> {
205 self.leaves
206 .iter()
207 .map(|(leaf_index, leaf)| (LeafIndex::new_max_depth(*leaf_index), leaf))
208 }
209
210 /// Returns an iterator over the key-value pairs of this [Smt].
211 pub fn entries(&self) -> impl Iterator<Item = &(RpoDigest, Word)> {
212 self.leaves().flat_map(|(_, leaf)| leaf.entries())
213 }
214
215 /// Returns an iterator over the inner nodes of this [Smt].
216 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
217 self.inner_nodes.values().map(|e| InnerNodeInfo {
218 value: e.hash(),
219 left: e.left,
220 right: e.right,
221 })
222 }
223
224 // STATE MUTATORS
225 // --------------------------------------------------------------------------------------------
226
227 /// Inserts a value at the specified key, returning the previous value associated with that key.
228 /// Recall that by definition, any key that hasn't been updated is associated with
229 /// [`Self::EMPTY_VALUE`].
230 ///
231 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
232 /// updating the root itself.
233 pub fn insert(&mut self, key: RpoDigest, value: Word) -> Word {
234 <Self as SparseMerkleTree<SMT_DEPTH>>::insert(self, key, value)
235 }
236
237 /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
238 /// tree, allowing for validation before applying those changes.
239 ///
240 /// This method returns a [`MutationSet`], which contains all the information for inserting
241 /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
242 /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
243 /// [`Smt::apply_mutations()`] can be called in order to commit these changes to the Merkle
244 /// tree, or [`drop()`] to discard them.
245 ///
246 /// # Example
247 /// ```
248 /// # use miden_crypto::{hash::rpo::RpoDigest, Felt, Word};
249 /// # use miden_crypto::merkle::{Smt, EmptySubtreeRoots, SMT_DEPTH};
250 /// let mut smt = Smt::new();
251 /// let pair = (RpoDigest::default(), Word::default());
252 /// let mutations = smt.compute_mutations(vec![pair]);
253 /// assert_eq!(mutations.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
254 /// smt.apply_mutations(mutations);
255 /// assert_eq!(smt.root(), *EmptySubtreeRoots::entry(SMT_DEPTH, 0));
256 /// ```
257 pub fn compute_mutations(
258 &self,
259 kv_pairs: impl IntoIterator<Item = (RpoDigest, Word)>,
260 ) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
261 <Self as SparseMerkleTree<SMT_DEPTH>>::compute_mutations(self, kv_pairs)
262 }
263
264 /// Applies the prospective mutations computed with [`Smt::compute_mutations()`] to this tree.
265 ///
266 /// # Errors
267 /// If `mutations` was computed on a tree with a different root than this one, returns
268 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
269 /// the `mutations` were computed against, and the second item is the actual current root of
270 /// this tree.
271 pub fn apply_mutations(
272 &mut self,
273 mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
274 ) -> Result<(), MerkleError> {
275 <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations(self, mutations)
276 }
277
278 /// Applies the prospective mutations computed with [`Smt::compute_mutations()`] to this tree
279 /// and returns the reverse mutation set.
280 ///
281 /// Applying the reverse mutation sets to the updated tree will revert the changes.
282 ///
283 /// # Errors
284 /// If `mutations` was computed on a tree with a different root than this one, returns
285 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
286 /// the `mutations` were computed against, and the second item is the actual current root of
287 /// this tree.
288 pub fn apply_mutations_with_reversion(
289 &mut self,
290 mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
291 ) -> Result<MutationSet<SMT_DEPTH, RpoDigest, Word>, MerkleError> {
292 <Self as SparseMerkleTree<SMT_DEPTH>>::apply_mutations_with_reversion(self, mutations)
293 }
294
295 // HELPERS
296 // --------------------------------------------------------------------------------------------
297
298 /// Inserts `value` at leaf index pointed to by `key`. `value` is guaranteed to not be the empty
299 /// value, such that this is indeed an insertion.
300 fn perform_insert(&mut self, key: RpoDigest, value: Word) -> Option<Word> {
301 debug_assert_ne!(value, Self::EMPTY_VALUE);
302
303 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
304
305 match self.leaves.get_mut(&leaf_index.value()) {
306 Some(leaf) => leaf.insert(key, value),
307 None => {
308 self.leaves.insert(leaf_index.value(), SmtLeaf::Single((key, value)));
309
310 None
311 },
312 }
313 }
314
315 /// Removes key-value pair at leaf index pointed to by `key` if it exists.
316 fn perform_remove(&mut self, key: RpoDigest) -> Option<Word> {
317 let leaf_index: LeafIndex<SMT_DEPTH> = Self::key_to_leaf_index(&key);
318
319 if let Some(leaf) = self.leaves.get_mut(&leaf_index.value()) {
320 let (old_value, is_empty) = leaf.remove(key);
321 if is_empty {
322 self.leaves.remove(&leaf_index.value());
323 }
324 old_value
325 } else {
326 // there's nothing stored at the leaf; nothing to update
327 None
328 }
329 }
330}
331
332impl SparseMerkleTree<SMT_DEPTH> for Smt {
333 type Key = RpoDigest;
334 type Value = Word;
335 type Leaf = SmtLeaf;
336 type Opening = SmtProof;
337
338 const EMPTY_VALUE: Self::Value = EMPTY_WORD;
339 const EMPTY_ROOT: RpoDigest = *EmptySubtreeRoots::entry(SMT_DEPTH, 0);
340
341 fn from_raw_parts(
342 inner_nodes: BTreeMap<NodeIndex, InnerNode>,
343 leaves: BTreeMap<u64, SmtLeaf>,
344 root: RpoDigest,
345 ) -> Result<Self, MerkleError> {
346 if cfg!(debug_assertions) {
347 let root_node = inner_nodes.get(&NodeIndex::root()).unwrap();
348 assert_eq!(root_node.hash(), root);
349 }
350
351 Ok(Self { root, inner_nodes, leaves })
352 }
353
354 fn root(&self) -> RpoDigest {
355 self.root
356 }
357
358 fn set_root(&mut self, root: RpoDigest) {
359 self.root = root;
360 }
361
362 fn get_inner_node(&self, index: NodeIndex) -> InnerNode {
363 self.inner_nodes
364 .get(&index)
365 .cloned()
366 .unwrap_or_else(|| EmptySubtreeRoots::get_inner_node(SMT_DEPTH, index.depth()))
367 }
368
369 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) -> Option<InnerNode> {
370 self.inner_nodes.insert(index, inner_node)
371 }
372
373 fn remove_inner_node(&mut self, index: NodeIndex) -> Option<InnerNode> {
374 self.inner_nodes.remove(&index)
375 }
376
377 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value> {
378 // inserting an `EMPTY_VALUE` is equivalent to removing any value associated with `key`
379 if value != Self::EMPTY_VALUE {
380 self.perform_insert(key, value)
381 } else {
382 self.perform_remove(key)
383 }
384 }
385
386 fn get_value(&self, key: &Self::Key) -> Self::Value {
387 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
388
389 match self.leaves.get(&leaf_pos) {
390 Some(leaf) => leaf.get_value(key).unwrap_or_default(),
391 None => EMPTY_WORD,
392 }
393 }
394
395 fn get_leaf(&self, key: &RpoDigest) -> Self::Leaf {
396 let leaf_pos = LeafIndex::<SMT_DEPTH>::from(*key).value();
397
398 match self.leaves.get(&leaf_pos) {
399 Some(leaf) => leaf.clone(),
400 None => SmtLeaf::new_empty(key.into()),
401 }
402 }
403
404 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest {
405 leaf.hash()
406 }
407
408 fn construct_prospective_leaf(
409 &self,
410 mut existing_leaf: SmtLeaf,
411 key: &RpoDigest,
412 value: &Word,
413 ) -> SmtLeaf {
414 debug_assert_eq!(existing_leaf.index(), Self::key_to_leaf_index(key));
415
416 match existing_leaf {
417 SmtLeaf::Empty(_) => SmtLeaf::new_single(*key, *value),
418 _ => {
419 if *value != EMPTY_WORD {
420 existing_leaf.insert(*key, *value);
421 } else {
422 existing_leaf.remove(*key);
423 }
424
425 existing_leaf
426 },
427 }
428 }
429
430 fn key_to_leaf_index(key: &RpoDigest) -> LeafIndex<SMT_DEPTH> {
431 let most_significant_felt = key[3];
432 LeafIndex::new_max_depth(most_significant_felt.as_int())
433 }
434
435 fn path_and_leaf_to_opening(path: MerklePath, leaf: SmtLeaf) -> SmtProof {
436 SmtProof::new_unchecked(path, leaf)
437 }
438
439 fn pairs_to_leaf(mut pairs: Vec<(RpoDigest, Word)>) -> SmtLeaf {
440 assert!(!pairs.is_empty());
441
442 if pairs.len() > 1 {
443 SmtLeaf::new_multiple(pairs).unwrap()
444 } else {
445 let (key, value) = pairs.pop().unwrap();
446 // TODO: should we ever be constructing empty leaves from pairs?
447 if value == Self::EMPTY_VALUE {
448 let index = Self::key_to_leaf_index(&key);
449 SmtLeaf::new_empty(index)
450 } else {
451 SmtLeaf::new_single(key, value)
452 }
453 }
454 }
455}
456
457impl Default for Smt {
458 fn default() -> Self {
459 Self::new()
460 }
461}
462
463// CONVERSIONS
464// ================================================================================================
465
466impl From<Word> for LeafIndex<SMT_DEPTH> {
467 fn from(value: Word) -> Self {
468 // We use the most significant `Felt` of a `Word` as the leaf index.
469 Self::new_max_depth(value[3].as_int())
470 }
471}
472
473impl From<RpoDigest> for LeafIndex<SMT_DEPTH> {
474 fn from(value: RpoDigest) -> Self {
475 Word::from(value).into()
476 }
477}
478
479impl From<&RpoDigest> for LeafIndex<SMT_DEPTH> {
480 fn from(value: &RpoDigest) -> Self {
481 Word::from(value).into()
482 }
483}
484
485// SERIALIZATION
486// ================================================================================================
487
488impl Serializable for Smt {
489 fn write_into<W: ByteWriter>(&self, target: &mut W) {
490 // Write the number of filled leaves for this Smt
491 target.write_usize(self.entries().count());
492
493 // Write each (key, value) pair
494 for (key, value) in self.entries() {
495 target.write(key);
496 target.write(value);
497 }
498 }
499
500 fn get_size_hint(&self) -> usize {
501 let entries_count = self.entries().count();
502
503 // Each entry is the size of a digest plus a word.
504 entries_count.get_size_hint()
505 + entries_count * (RpoDigest::SERIALIZED_SIZE + EMPTY_WORD.get_size_hint())
506 }
507}
508
509impl Deserializable for Smt {
510 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
511 // Read the number of filled leaves for this Smt
512 let num_filled_leaves = source.read_usize()?;
513 let mut entries = Vec::with_capacity(num_filled_leaves);
514
515 for _ in 0..num_filled_leaves {
516 let key = source.read()?;
517 let value = source.read()?;
518 entries.push((key, value));
519 }
520
521 Self::with_entries(entries)
522 .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
523 }
524}
525
526#[test]
527fn test_smt_serialization_deserialization() {
528 // Smt for default types (empty map)
529 let smt_default = Smt::default();
530 let bytes = smt_default.to_bytes();
531 assert_eq!(smt_default, Smt::read_from_bytes(&bytes).unwrap());
532 assert_eq!(bytes.len(), smt_default.get_size_hint());
533
534 // Smt with values
535 let smt_leaves_2: [(RpoDigest, Word); 2] = [
536 (
537 RpoDigest::new([Felt::new(101), Felt::new(102), Felt::new(103), Felt::new(104)]),
538 [Felt::new(1_u64), Felt::new(2_u64), Felt::new(3_u64), Felt::new(4_u64)],
539 ),
540 (
541 RpoDigest::new([Felt::new(105), Felt::new(106), Felt::new(107), Felt::new(108)]),
542 [Felt::new(5_u64), Felt::new(6_u64), Felt::new(7_u64), Felt::new(8_u64)],
543 ),
544 ];
545 let smt = Smt::with_entries(smt_leaves_2).unwrap();
546
547 let bytes = smt.to_bytes();
548 assert_eq!(smt, Smt::read_from_bytes(&bytes).unwrap());
549 assert_eq!(bytes.len(), smt.get_size_hint());
550}