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