miden-crypto/src/merkle/smt/mod.rs
src/merkle/smt/mod.rsRaw
1use alloc::{collections::BTreeMap, vec::Vec};
2use core::{hash::Hash, mem};
3
4use num::Integer;
5use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
6
7use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
8use crate::{
9 hash::rpo::{Rpo256, RpoDigest},
10 Felt, Word, EMPTY_WORD,
11};
12
13mod full;
14pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
15
16mod simple;
17pub use simple::SimpleSmt;
18
19// CONSTANTS
20// ================================================================================================
21
22/// Minimum supported depth.
23pub const SMT_MIN_DEPTH: u8 = 1;
24
25/// Maximum supported depth.
26pub const SMT_MAX_DEPTH: u8 = 64;
27
28// SPARSE MERKLE TREE
29// ================================================================================================
30
31/// A map whose keys are not guarantied to be ordered.
32#[cfg(feature = "smt_hashmaps")]
33type UnorderedMap<K, V> = hashbrown::HashMap<K, V>;
34#[cfg(not(feature = "smt_hashmaps"))]
35type UnorderedMap<K, V> = alloc::collections::BTreeMap<K, V>;
36type InnerNodes = UnorderedMap<NodeIndex, InnerNode>;
37type Leaves<T> = UnorderedMap<u64, T>;
38type NodeMutations = UnorderedMap<NodeIndex, NodeMutation>;
39
40/// An abstract description of a sparse Merkle tree.
41///
42/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
43/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
44/// value was not explicitly set, then its value is the default value. Typically, the vast majority
45/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
46/// representation of the tree will only keep track of the leaves that have a different value from
47/// the default.
48///
49/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
50/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
51/// tree is just a single value, and is probably a programming mistake.
52///
53/// Every key maps to one leaf. If there are as many keys as there are leaves, then
54/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
55/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
56/// must accommodate all keys that map to the same leaf.
57///
58/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
59pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
60 /// The type for a key
61 type Key: Clone + Ord + Eq + Hash;
62 /// The type for a value
63 type Value: Clone + PartialEq;
64 /// The type for a leaf
65 type Leaf: Clone;
66 /// The type for an opening (i.e. a "proof") of a leaf
67 type Opening;
68
69 /// The default value used to compute the hash of empty leaves
70 const EMPTY_VALUE: Self::Value;
71
72 /// The root of the empty tree with provided DEPTH
73 const EMPTY_ROOT: RpoDigest;
74
75 // PROVIDED METHODS
76 // ---------------------------------------------------------------------------------------------
77
78 /// Creates a new sparse Merkle tree from an existing set of key-value pairs, in parallel.
79 #[cfg(feature = "concurrent")]
80 fn with_entries_par(entries: Vec<(Self::Key, Self::Value)>) -> Result<Self, MerkleError>
81 where
82 Self: Sized,
83 {
84 let (inner_nodes, leaves) = Self::build_subtrees(entries);
85 let root = inner_nodes.get(&NodeIndex::root()).unwrap().hash();
86 Self::from_raw_parts(inner_nodes, leaves, root)
87 }
88
89 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
90 /// path to the leaf, as well as the leaf itself.
91 fn open(&self, key: &Self::Key) -> Self::Opening {
92 let leaf = self.get_leaf(key);
93
94 let mut index: NodeIndex = {
95 let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
96 leaf_index.into()
97 };
98
99 let merkle_path = {
100 let mut path = Vec::with_capacity(index.depth() as usize);
101 for _ in 0..index.depth() {
102 let is_right = index.is_value_odd();
103 index.move_up();
104 let InnerNode { left, right } = self.get_inner_node(index);
105 let value = if is_right { left } else { right };
106 path.push(value);
107 }
108
109 MerklePath::new(path)
110 };
111
112 Self::path_and_leaf_to_opening(merkle_path, leaf)
113 }
114
115 /// Inserts a value at the specified key, returning the previous value associated with that key.
116 /// Recall that by definition, any key that hasn't been updated is associated with
117 /// [`Self::EMPTY_VALUE`].
118 ///
119 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
120 /// updating the root itself.
121 fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
122 let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
123
124 // if the old value and new value are the same, there is nothing to update
125 if value == old_value {
126 return value;
127 }
128
129 let leaf = self.get_leaf(&key);
130 let node_index = {
131 let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
132 leaf_index.into()
133 };
134
135 self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
136
137 old_value
138 }
139
140 /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
141 /// `node_hash_at_index` is the hash of the node stored at index.
142 fn recompute_nodes_from_index_to_root(
143 &mut self,
144 mut index: NodeIndex,
145 node_hash_at_index: RpoDigest,
146 ) {
147 let mut node_hash = node_hash_at_index;
148 for node_depth in (0..index.depth()).rev() {
149 let is_right = index.is_value_odd();
150 index.move_up();
151 let InnerNode { left, right } = self.get_inner_node(index);
152 let (left, right) = if is_right {
153 (left, node_hash)
154 } else {
155 (node_hash, right)
156 };
157 node_hash = Rpo256::merge(&[left, right]);
158
159 if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
160 // If a subtree is empty, then can remove the inner node, since it's equal to the
161 // default value
162 self.remove_inner_node(index);
163 } else {
164 self.insert_inner_node(index, InnerNode { left, right });
165 }
166 }
167 self.set_root(node_hash);
168 }
169
170 /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
171 /// tree, allowing for validation before applying those changes.
172 ///
173 /// This method returns a [`MutationSet`], which contains all the information for inserting
174 /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
175 /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
176 /// [`SparseMerkleTree::apply_mutations()`] can be called in order to commit these changes to
177 /// the Merkle tree, or [`drop()`] to discard them.
178 fn compute_mutations(
179 &self,
180 kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
181 ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
182 use NodeMutation::*;
183
184 let mut new_root = self.root();
185 let mut new_pairs: UnorderedMap<Self::Key, Self::Value> = Default::default();
186 let mut node_mutations: NodeMutations = Default::default();
187
188 for (key, value) in kv_pairs {
189 // If the old value and the new value are the same, there is nothing to update.
190 // For the unusual case that kv_pairs has multiple values at the same key, we'll have
191 // to check the key-value pairs we've already seen to get the "effective" old value.
192 let old_value = new_pairs.get(&key).cloned().unwrap_or_else(|| self.get_value(&key));
193 if value == old_value {
194 continue;
195 }
196
197 let leaf_index = Self::key_to_leaf_index(&key);
198 let mut node_index = NodeIndex::from(leaf_index);
199
200 // We need the current leaf's hash to calculate the new leaf, but in the rare case that
201 // `kv_pairs` has multiple pairs that go into the same leaf, then those pairs are also
202 // part of the "current leaf".
203 let old_leaf = {
204 let pairs_at_index = new_pairs
205 .iter()
206 .filter(|&(new_key, _)| Self::key_to_leaf_index(new_key) == leaf_index);
207
208 pairs_at_index.fold(self.get_leaf(&key), |acc, (k, v)| {
209 // Most of the time `pairs_at_index` should only contain a single entry (or
210 // none at all), as multi-leaves should be really rare.
211 let existing_leaf = acc.clone();
212 self.construct_prospective_leaf(existing_leaf, k, v)
213 })
214 };
215
216 let new_leaf = self.construct_prospective_leaf(old_leaf, &key, &value);
217
218 let mut new_child_hash = Self::hash_leaf(&new_leaf);
219
220 for node_depth in (0..node_index.depth()).rev() {
221 // Whether the node we're replacing is the right child or the left child.
222 let is_right = node_index.is_value_odd();
223 node_index.move_up();
224
225 let old_node = node_mutations
226 .get(&node_index)
227 .map(|mutation| match mutation {
228 Addition(node) => node.clone(),
229 Removal => EmptySubtreeRoots::get_inner_node(DEPTH, node_depth),
230 })
231 .unwrap_or_else(|| self.get_inner_node(node_index));
232
233 let new_node = if is_right {
234 InnerNode {
235 left: old_node.left,
236 right: new_child_hash,
237 }
238 } else {
239 InnerNode {
240 left: new_child_hash,
241 right: old_node.right,
242 }
243 };
244
245 // The next iteration will operate on this new node's hash.
246 new_child_hash = new_node.hash();
247
248 let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
249 let is_removal = new_child_hash == equivalent_empty_hash;
250 let new_entry = if is_removal { Removal } else { Addition(new_node) };
251 node_mutations.insert(node_index, new_entry);
252 }
253
254 // Once we're at depth 0, the last node we made is the new root.
255 new_root = new_child_hash;
256 // And then we're done with this pair; on to the next one.
257 new_pairs.insert(key, value);
258 }
259
260 MutationSet {
261 old_root: self.root(),
262 new_root,
263 node_mutations,
264 new_pairs,
265 }
266 }
267
268 /// Applies the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
269 /// this tree.
270 ///
271 /// # Errors
272 /// If `mutations` was computed on a tree with a different root than this one, returns
273 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
274 /// the `mutations` were computed against, and the second item is the actual current root of
275 /// this tree.
276 fn apply_mutations(
277 &mut self,
278 mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
279 ) -> Result<(), MerkleError>
280 where
281 Self: Sized,
282 {
283 use NodeMutation::*;
284 let MutationSet {
285 old_root,
286 node_mutations,
287 new_pairs,
288 new_root,
289 } = mutations;
290
291 // Guard against accidentally trying to apply mutations that were computed against a
292 // different tree, including a stale version of this tree.
293 if old_root != self.root() {
294 return Err(MerkleError::ConflictingRoots {
295 expected_root: self.root(),
296 actual_root: old_root,
297 });
298 }
299
300 for (index, mutation) in node_mutations {
301 match mutation {
302 Removal => {
303 self.remove_inner_node(index);
304 },
305 Addition(node) => {
306 self.insert_inner_node(index, node);
307 },
308 }
309 }
310
311 for (key, value) in new_pairs {
312 self.insert_value(key, value);
313 }
314
315 self.set_root(new_root);
316
317 Ok(())
318 }
319
320 /// Applies the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
321 /// this tree and returns the reverse mutation set. Applying the reverse mutation sets to the
322 /// updated tree will revert the changes.
323 ///
324 /// # Errors
325 /// If `mutations` was computed on a tree with a different root than this one, returns
326 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
327 /// the `mutations` were computed against, and the second item is the actual current root of
328 /// this tree.
329 fn apply_mutations_with_reversion(
330 &mut self,
331 mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
332 ) -> Result<MutationSet<DEPTH, Self::Key, Self::Value>, MerkleError>
333 where
334 Self: Sized,
335 {
336 use NodeMutation::*;
337 let MutationSet {
338 old_root,
339 node_mutations,
340 new_pairs,
341 new_root,
342 } = mutations;
343
344 // Guard against accidentally trying to apply mutations that were computed against a
345 // different tree, including a stale version of this tree.
346 if old_root != self.root() {
347 return Err(MerkleError::ConflictingRoots {
348 expected_root: self.root(),
349 actual_root: old_root,
350 });
351 }
352
353 let mut reverse_mutations = NodeMutations::new();
354 for (index, mutation) in node_mutations {
355 match mutation {
356 Removal => {
357 if let Some(node) = self.remove_inner_node(index) {
358 reverse_mutations.insert(index, Addition(node));
359 }
360 },
361 Addition(node) => {
362 if let Some(old_node) = self.insert_inner_node(index, node) {
363 reverse_mutations.insert(index, Addition(old_node));
364 } else {
365 reverse_mutations.insert(index, Removal);
366 }
367 },
368 }
369 }
370
371 let mut reverse_pairs = UnorderedMap::new();
372 for (key, value) in new_pairs {
373 if let Some(old_value) = self.insert_value(key.clone(), value) {
374 reverse_pairs.insert(key, old_value);
375 } else {
376 reverse_pairs.insert(key, Self::EMPTY_VALUE);
377 }
378 }
379
380 self.set_root(new_root);
381
382 Ok(MutationSet {
383 old_root: new_root,
384 node_mutations: reverse_mutations,
385 new_pairs: reverse_pairs,
386 new_root: old_root,
387 })
388 }
389
390 // REQUIRED METHODS
391 // ---------------------------------------------------------------------------------------------
392
393 /// Construct this type from already computed leaves and nodes. The caller ensures passed
394 /// arguments are correct and consistent with each other.
395 fn from_raw_parts(
396 inner_nodes: InnerNodes,
397 leaves: Leaves<Self::Leaf>,
398 root: RpoDigest,
399 ) -> Result<Self, MerkleError>
400 where
401 Self: Sized;
402
403 /// The root of the tree
404 fn root(&self) -> RpoDigest;
405
406 /// Sets the root of the tree
407 fn set_root(&mut self, root: RpoDigest);
408
409 /// Retrieves an inner node at the given index
410 fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
411
412 /// Inserts an inner node at the given index
413 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) -> Option<InnerNode>;
414
415 /// Removes an inner node at the given index
416 fn remove_inner_node(&mut self, index: NodeIndex) -> Option<InnerNode>;
417
418 /// Inserts a leaf node, and returns the value at the key if already exists
419 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
420
421 /// Returns the value at the specified key. Recall that by definition, any key that hasn't been
422 /// updated is associated with [`Self::EMPTY_VALUE`].
423 fn get_value(&self, key: &Self::Key) -> Self::Value;
424
425 /// Returns the leaf at the specified index.
426 fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
427
428 /// Returns the hash of a leaf
429 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
430
431 /// Returns what a leaf would look like if a key-value pair were inserted into the tree, without
432 /// mutating the tree itself. The existing leaf can be empty.
433 ///
434 /// To get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)`
435 /// as the argument for `existing_leaf`. The return value from this function can be chained back
436 /// into this function as the first argument to continue making prospective changes.
437 ///
438 /// # Invariants
439 /// Because this method is for a prospective key-value insertion into a specific leaf,
440 /// `existing_leaf` must have the same leaf index as `key` (as determined by
441 /// [`SparseMerkleTree::key_to_leaf_index()`]), or the result will be meaningless.
442 fn construct_prospective_leaf(
443 &self,
444 existing_leaf: Self::Leaf,
445 key: &Self::Key,
446 value: &Self::Value,
447 ) -> Self::Leaf;
448
449 /// Maps a key to a leaf index
450 fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
451
452 /// Constructs a single leaf from an arbitrary amount of key-value pairs.
453 /// Those pairs must all have the same leaf index.
454 fn pairs_to_leaf(pairs: Vec<(Self::Key, Self::Value)>) -> Self::Leaf;
455
456 /// Maps a (MerklePath, Self::Leaf) to an opening.
457 ///
458 /// The length `path` is guaranteed to be equal to `DEPTH`
459 fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
460
461 /// Performs the initial transforms for constructing a [`SparseMerkleTree`] by composing
462 /// subtrees. In other words, this function takes the key-value inputs to the tree, and produces
463 /// the inputs to feed into [`build_subtree()`].
464 ///
465 /// `pairs` *must* already be sorted **by leaf index column**, not simply sorted by key. If
466 /// `pairs` is not correctly sorted, the returned computations will be incorrect.
467 ///
468 /// # Panics
469 /// With debug assertions on, this function panics if it detects that `pairs` is not correctly
470 /// sorted. Without debug assertions, the returned computations will be incorrect.
471 fn sorted_pairs_to_leaves(
472 pairs: Vec<(Self::Key, Self::Value)>,
473 ) -> PairComputations<u64, Self::Leaf> {
474 debug_assert!(pairs.is_sorted_by_key(|(key, _)| Self::key_to_leaf_index(key).value()));
475
476 let mut accumulator: PairComputations<u64, Self::Leaf> = Default::default();
477 let mut accumulated_leaves: Vec<SubtreeLeaf> = Vec::with_capacity(pairs.len() / 2);
478
479 // As we iterate, we'll keep track of the kv-pairs we've seen so far that correspond to a
480 // single leaf. When we see a pair that's in a different leaf, we'll swap these pairs
481 // out and store them in our accumulated leaves.
482 let mut current_leaf_buffer: Vec<(Self::Key, Self::Value)> = Default::default();
483
484 let mut iter = pairs.into_iter().peekable();
485 while let Some((key, value)) = iter.next() {
486 let col = Self::key_to_leaf_index(&key).index.value();
487 let peeked_col = iter.peek().map(|(key, _v)| {
488 let index = Self::key_to_leaf_index(key);
489 let next_col = index.index.value();
490 // We panic if `pairs` is not sorted by column.
491 debug_assert!(next_col >= col);
492 next_col
493 });
494 current_leaf_buffer.push((key, value));
495
496 // If the next pair is the same column as this one, then we're done after adding this
497 // pair to the buffer.
498 if peeked_col == Some(col) {
499 continue;
500 }
501
502 // Otherwise, the next pair is a different column, or there is no next pair. Either way
503 // it's time to swap out our buffer.
504 let leaf_pairs = mem::take(&mut current_leaf_buffer);
505 let leaf = Self::pairs_to_leaf(leaf_pairs);
506 let hash = Self::hash_leaf(&leaf);
507
508 accumulator.nodes.insert(col, leaf);
509 accumulated_leaves.push(SubtreeLeaf { col, hash });
510
511 debug_assert!(current_leaf_buffer.is_empty());
512 }
513
514 // TODO: determine is there is any notable performance difference between computing
515 // subtree boundaries after the fact as an iterator adapter (like this), versus computing
516 // subtree boundaries as we go. Either way this function is only used at the beginning of a
517 // parallel construction, so it should not be a critical path.
518 accumulator.leaves = SubtreeLeavesIter::from_leaves(&mut accumulated_leaves).collect();
519 accumulator
520 }
521
522 /// Computes the raw parts for a new sparse Merkle tree from a set of key-value pairs.
523 ///
524 /// `entries` need not be sorted. This function will sort them.
525 #[cfg(feature = "concurrent")]
526 fn build_subtrees(
527 mut entries: Vec<(Self::Key, Self::Value)>,
528 ) -> (InnerNodes, Leaves<Self::Leaf>) {
529 entries.sort_by_key(|item| {
530 let index = Self::key_to_leaf_index(&item.0);
531 index.value()
532 });
533 Self::build_subtrees_from_sorted_entries(entries)
534 }
535
536 /// Computes the raw parts for a new sparse Merkle tree from a set of key-value pairs.
537 ///
538 /// This function is mostly an implementation detail of
539 /// [`SparseMerkleTree::with_entries_par()`].
540 #[cfg(feature = "concurrent")]
541 fn build_subtrees_from_sorted_entries(
542 entries: Vec<(Self::Key, Self::Value)>,
543 ) -> (InnerNodes, Leaves<Self::Leaf>) {
544 use rayon::prelude::*;
545
546 let mut accumulated_nodes: InnerNodes = Default::default();
547
548 let PairComputations {
549 leaves: mut leaf_subtrees,
550 nodes: initial_leaves,
551 } = Self::sorted_pairs_to_leaves(entries);
552
553 for current_depth in (SUBTREE_DEPTH..=DEPTH).step_by(SUBTREE_DEPTH as usize).rev() {
554 let (nodes, mut subtree_roots): (Vec<BTreeMap<_, _>>, Vec<SubtreeLeaf>) = leaf_subtrees
555 .into_par_iter()
556 .map(|subtree| {
557 debug_assert!(subtree.is_sorted());
558 debug_assert!(!subtree.is_empty());
559
560 let (nodes, subtree_root) = build_subtree(subtree, DEPTH, current_depth);
561 (nodes, subtree_root)
562 })
563 .unzip();
564
565 leaf_subtrees = SubtreeLeavesIter::from_leaves(&mut subtree_roots).collect();
566 accumulated_nodes.extend(nodes.into_iter().flatten());
567
568 debug_assert!(!leaf_subtrees.is_empty());
569 }
570 (accumulated_nodes, initial_leaves)
571 }
572}
573
574// INNER NODE
575// ================================================================================================
576
577/// This struct is public so functions returning it can be used in `benches/`, but is otherwise not
578/// part of the public API.
579#[doc(hidden)]
580#[derive(Debug, Default, Clone, PartialEq, Eq)]
581#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
582pub struct InnerNode {
583 pub left: RpoDigest,
584 pub right: RpoDigest,
585}
586
587impl InnerNode {
588 pub fn hash(&self) -> RpoDigest {
589 Rpo256::merge(&[self.left, self.right])
590 }
591}
592
593// LEAF INDEX
594// ================================================================================================
595
596/// The index of a leaf, at a depth known at compile-time.
597#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
598#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
599pub struct LeafIndex<const DEPTH: u8> {
600 index: NodeIndex,
601}
602
603impl<const DEPTH: u8> LeafIndex<DEPTH> {
604 pub fn new(value: u64) -> Result<Self, MerkleError> {
605 if DEPTH < SMT_MIN_DEPTH {
606 return Err(MerkleError::DepthTooSmall(DEPTH));
607 }
608
609 Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
610 }
611
612 pub fn value(&self) -> u64 {
613 self.index.value()
614 }
615}
616
617impl LeafIndex<SMT_MAX_DEPTH> {
618 pub const fn new_max_depth(value: u64) -> Self {
619 LeafIndex {
620 index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
621 }
622 }
623}
624
625impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
626 fn from(value: LeafIndex<DEPTH>) -> Self {
627 value.index
628 }
629}
630
631impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
632 type Error = MerkleError;
633
634 fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
635 if node_index.depth() != DEPTH {
636 return Err(MerkleError::InvalidNodeIndexDepth {
637 expected: DEPTH,
638 provided: node_index.depth(),
639 });
640 }
641
642 Self::new(node_index.value())
643 }
644}
645
646// MUTATIONS
647// ================================================================================================
648
649/// A change to an inner node of a [`SparseMerkleTree`] that hasn't yet been applied.
650/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
651/// need to occur at which node indices.
652#[derive(Debug, Clone, PartialEq, Eq)]
653pub enum NodeMutation {
654 /// Corresponds to [`SparseMerkleTree::remove_inner_node()`].
655 Removal,
656 /// Corresponds to [`SparseMerkleTree::insert_inner_node()`].
657 Addition(InnerNode),
658}
659
660/// Represents a group of prospective mutations to a `SparseMerkleTree`, created by
661/// `SparseMerkleTree::compute_mutations()`, and that can be applied with
662/// `SparseMerkleTree::apply_mutations()`.
663#[derive(Debug, Clone, Default, PartialEq, Eq)]
664pub struct MutationSet<const DEPTH: u8, K: Eq + Hash, V> {
665 /// The root of the Merkle tree this MutationSet is for, recorded at the time
666 /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying
667 /// mutations to the wrong tree or applying stale mutations to a tree that has since changed.
668 old_root: RpoDigest,
669 /// The set of nodes that need to be removed or added. The "effective" node at an index is the
670 /// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that
671 /// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a
672 /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`]
673 /// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call.
674 node_mutations: NodeMutations,
675 /// The set of top-level key-value pairs we're prospectively adding to the tree, including
676 /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling
677 /// back to the existing value in the Merkle tree. Each entry corresponds to a
678 /// [`SparseMerkleTree::insert_value()`] call.
679 new_pairs: UnorderedMap<K, V>,
680 /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with
681 /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call.
682 new_root: RpoDigest,
683}
684
685impl<const DEPTH: u8, K: Eq + Hash, V> MutationSet<DEPTH, K, V> {
686 /// Returns the SMT root that was calculated during `SparseMerkleTree::compute_mutations()`. See
687 /// that method for more information.
688 pub fn root(&self) -> RpoDigest {
689 self.new_root
690 }
691
692 /// Returns the SMT root before the mutations were applied.
693 pub fn old_root(&self) -> RpoDigest {
694 self.old_root
695 }
696
697 /// Returns the set of inner nodes that need to be removed or added.
698 pub fn node_mutations(&self) -> &NodeMutations {
699 &self.node_mutations
700 }
701
702 /// Returns the set of top-level key-value pairs that need to be added, updated or deleted
703 /// (i.e. set to `EMPTY_WORD`).
704 pub fn new_pairs(&self) -> &UnorderedMap<K, V> {
705 &self.new_pairs
706 }
707}
708
709// SERIALIZATION
710// ================================================================================================
711
712impl Serializable for InnerNode {
713 fn write_into<W: ByteWriter>(&self, target: &mut W) {
714 target.write(self.left);
715 target.write(self.right);
716 }
717}
718
719impl Deserializable for InnerNode {
720 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
721 let left = source.read()?;
722 let right = source.read()?;
723
724 Ok(Self { left, right })
725 }
726}
727
728impl Serializable for NodeMutation {
729 fn write_into<W: ByteWriter>(&self, target: &mut W) {
730 match self {
731 NodeMutation::Removal => target.write_bool(false),
732 NodeMutation::Addition(inner_node) => {
733 target.write_bool(true);
734 inner_node.write_into(target);
735 },
736 }
737 }
738}
739
740impl Deserializable for NodeMutation {
741 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
742 if source.read_bool()? {
743 let inner_node = source.read()?;
744 return Ok(NodeMutation::Addition(inner_node));
745 }
746
747 Ok(NodeMutation::Removal)
748 }
749}
750
751impl<const DEPTH: u8, K: Serializable + Eq + Hash, V: Serializable> Serializable
752 for MutationSet<DEPTH, K, V>
753{
754 fn write_into<W: ByteWriter>(&self, target: &mut W) {
755 target.write(self.old_root);
756 target.write(self.new_root);
757
758 let inner_removals: Vec<_> = self
759 .node_mutations
760 .iter()
761 .filter(|(_, value)| matches!(value, NodeMutation::Removal))
762 .map(|(key, _)| key)
763 .collect();
764 let inner_additions: Vec<_> = self
765 .node_mutations
766 .iter()
767 .filter_map(|(key, value)| match value {
768 NodeMutation::Addition(node) => Some((key, node)),
769 _ => None,
770 })
771 .collect();
772
773 target.write(inner_removals);
774 target.write(inner_additions);
775
776 target.write_usize(self.new_pairs.len());
777 target.write_many(&self.new_pairs);
778 }
779}
780
781impl<const DEPTH: u8, K: Deserializable + Ord + Eq + Hash, V: Deserializable> Deserializable
782 for MutationSet<DEPTH, K, V>
783{
784 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
785 let old_root = source.read()?;
786 let new_root = source.read()?;
787
788 let inner_removals: Vec<NodeIndex> = source.read()?;
789 let inner_additions: Vec<(NodeIndex, InnerNode)> = source.read()?;
790
791 let node_mutations = NodeMutations::from_iter(
792 inner_removals.into_iter().map(|index| (index, NodeMutation::Removal)).chain(
793 inner_additions
794 .into_iter()
795 .map(|(index, node)| (index, NodeMutation::Addition(node))),
796 ),
797 );
798
799 let num_new_pairs = source.read_usize()?;
800 let new_pairs = source.read_many(num_new_pairs)?;
801 let new_pairs = UnorderedMap::from_iter(new_pairs);
802
803 Ok(Self {
804 old_root,
805 node_mutations,
806 new_pairs,
807 new_root,
808 })
809 }
810}
811
812// SUBTREES
813// ================================================================================================
814
815/// A subtree is of depth 8.
816const SUBTREE_DEPTH: u8 = 8;
817
818/// A depth-8 subtree contains 256 "columns" that can possibly be occupied.
819const COLS_PER_SUBTREE: u64 = u64::pow(2, SUBTREE_DEPTH as u32);
820
821/// Helper struct for organizing the data we care about when computing Merkle subtrees.
822///
823/// Note that these represet "conceptual" leaves of some subtree, not necessarily
824/// the leaf type for the sparse Merkle tree.
825#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default)]
826pub struct SubtreeLeaf {
827 /// The 'value' field of [`NodeIndex`]. When computing a subtree, the depth is already known.
828 pub col: u64,
829 /// The hash of the node this `SubtreeLeaf` represents.
830 pub hash: RpoDigest,
831}
832
833/// Helper struct to organize the return value of [`SparseMerkleTree::sorted_pairs_to_leaves()`].
834#[derive(Debug, Clone)]
835pub(crate) struct PairComputations<K, L> {
836 /// Literal leaves to be added to the sparse Merkle tree's internal mapping.
837 pub nodes: UnorderedMap<K, L>,
838 /// "Conceptual" leaves that will be used for computations.
839 pub leaves: Vec<Vec<SubtreeLeaf>>,
840}
841
842// Derive requires `L` to impl Default, even though we don't actually need that.
843impl<K, L> Default for PairComputations<K, L> {
844 fn default() -> Self {
845 Self {
846 nodes: Default::default(),
847 leaves: Default::default(),
848 }
849 }
850}
851
852#[derive(Debug)]
853struct SubtreeLeavesIter<'s> {
854 leaves: core::iter::Peekable<alloc::vec::Drain<'s, SubtreeLeaf>>,
855}
856impl<'s> SubtreeLeavesIter<'s> {
857 fn from_leaves(leaves: &'s mut Vec<SubtreeLeaf>) -> Self {
858 // TODO: determine if there is any notable performance difference between taking a Vec,
859 // which many need flattening first, vs storing a `Box<dyn Iterator<Item = SubtreeLeaf>>`.
860 // The latter may have self-referential properties that are impossible to express in purely
861 // safe Rust Rust.
862 Self { leaves: leaves.drain(..).peekable() }
863 }
864}
865impl Iterator for SubtreeLeavesIter<'_> {
866 type Item = Vec<SubtreeLeaf>;
867
868 /// Each `next()` collects an entire subtree.
869 fn next(&mut self) -> Option<Vec<SubtreeLeaf>> {
870 let mut subtree: Vec<SubtreeLeaf> = Default::default();
871
872 let mut last_subtree_col = 0;
873
874 while let Some(leaf) = self.leaves.peek() {
875 last_subtree_col = u64::max(1, last_subtree_col);
876 let is_exact_multiple = Integer::is_multiple_of(&last_subtree_col, &COLS_PER_SUBTREE);
877 let next_subtree_col = if is_exact_multiple {
878 u64::next_multiple_of(last_subtree_col + 1, COLS_PER_SUBTREE)
879 } else {
880 last_subtree_col.next_multiple_of(COLS_PER_SUBTREE)
881 };
882
883 last_subtree_col = leaf.col;
884 if leaf.col < next_subtree_col {
885 subtree.push(self.leaves.next().unwrap());
886 } else if subtree.is_empty() {
887 continue;
888 } else {
889 break;
890 }
891 }
892
893 if subtree.is_empty() {
894 debug_assert!(self.leaves.peek().is_none());
895 return None;
896 }
897
898 Some(subtree)
899 }
900}
901
902// HELPER FUNCTIONS
903// ================================================================================================
904
905/// Builds Merkle nodes from a bottom layer of "leaves" -- represented by a horizontal index and
906/// the hash of the leaf at that index. `leaves` *must* be sorted by horizontal index, and
907/// `leaves` must not contain more than one depth-8 subtree's worth of leaves.
908///
909/// This function will then calculate the inner nodes above each leaf for 8 layers, as well as
910/// the "leaves" for the next 8-deep subtree, so this function can effectively be chained into
911/// itself.
912///
913/// # Panics
914/// With debug assertions on, this function panics under invalid inputs: if `leaves` contains
915/// more entries than can fit in a depth-8 subtree, if `leaves` contains leaves belonging to
916/// different depth-8 subtrees, if `bottom_depth` is lower in the tree than the specified
917/// maximum depth (`DEPTH`), or if `leaves` is not sorted.
918fn build_subtree(
919 mut leaves: Vec<SubtreeLeaf>,
920 tree_depth: u8,
921 bottom_depth: u8,
922) -> (BTreeMap<NodeIndex, InnerNode>, SubtreeLeaf) {
923 debug_assert!(bottom_depth <= tree_depth);
924 debug_assert!(Integer::is_multiple_of(&bottom_depth, &SUBTREE_DEPTH));
925 debug_assert!(leaves.len() <= usize::pow(2, SUBTREE_DEPTH as u32));
926 let subtree_root = bottom_depth - SUBTREE_DEPTH;
927 let mut inner_nodes: BTreeMap<NodeIndex, InnerNode> = Default::default();
928 let mut next_leaves: Vec<SubtreeLeaf> = Vec::with_capacity(leaves.len() / 2);
929 for next_depth in (subtree_root..bottom_depth).rev() {
930 debug_assert!(next_depth <= bottom_depth);
931 // `next_depth` is the stuff we're making.
932 // `current_depth` is the stuff we have.
933 let current_depth = next_depth + 1;
934 let mut iter = leaves.drain(..).peekable();
935 while let Some(first) = iter.next() {
936 // On non-continuous iterations, including the first iteration, `first_column` may
937 // be a left or right node. On subsequent continuous iterations, we will always call
938 // `iter.next()` twice.
939 // On non-continuous iterations (including the very first iteration), this column
940 // could be either on the left or the right. If the next iteration is not
941 // discontinuous with our right node, then the next iteration's
942 let is_right = first.col.is_odd();
943 let (left, right) = if is_right {
944 // Discontinuous iteration: we have no left node, so it must be empty.
945 let left = SubtreeLeaf {
946 col: first.col - 1,
947 hash: *EmptySubtreeRoots::entry(tree_depth, current_depth),
948 };
949 let right = first;
950 (left, right)
951 } else {
952 let left = first;
953 let right_col = first.col + 1;
954 let right = match iter.peek().copied() {
955 Some(SubtreeLeaf { col, .. }) if col == right_col => {
956 // Our inputs must be sorted.
957 debug_assert!(left.col <= col);
958 // The next leaf in the iterator is our sibling. Use it and consume it!
959 iter.next().unwrap()
960 },
961 // Otherwise, the leaves don't contain our sibling, so our sibling must be
962 // empty.
963 _ => SubtreeLeaf {
964 col: right_col,
965 hash: *EmptySubtreeRoots::entry(tree_depth, current_depth),
966 },
967 };
968 (left, right)
969 };
970 let index = NodeIndex::new_unchecked(current_depth, left.col).parent();
971 let node = InnerNode { left: left.hash, right: right.hash };
972 let hash = node.hash();
973 let &equivalent_empty_hash = EmptySubtreeRoots::entry(tree_depth, next_depth);
974 // If this hash is empty, then it doesn't become a new inner node, nor does it count
975 // as a leaf for the next depth.
976 if hash != equivalent_empty_hash {
977 inner_nodes.insert(index, node);
978 next_leaves.push(SubtreeLeaf { col: index.value(), hash });
979 }
980 }
981 // Stop borrowing `leaves`, so we can swap it.
982 // The iterator is empty at this point anyway.
983 drop(iter);
984 // After each depth, consider the stuff we just made the new "leaves", and empty the
985 // other collection.
986 mem::swap(&mut leaves, &mut next_leaves);
987 }
988 debug_assert_eq!(leaves.len(), 1);
989 let root = leaves.pop().unwrap();
990 (inner_nodes, root)
991}
992
993#[cfg(feature = "internal")]
994pub fn build_subtree_for_bench(
995 leaves: Vec<SubtreeLeaf>,
996 tree_depth: u8,
997 bottom_depth: u8,
998) -> (BTreeMap<NodeIndex, InnerNode>, SubtreeLeaf) {
999 build_subtree(leaves, tree_depth, bottom_depth)
1000}
1001
1002// TESTS
1003// ================================================================================================
1004#[cfg(test)]
1005mod tests;