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