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