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