miden-crypto/src/merkle/smt/mod.rs
src/merkle/smt/mod.rsRaw
1use alloc::{collections::BTreeMap, vec::Vec};
2
3use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
4
5use super::{EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex};
6use crate::{
7 hash::rpo::{Rpo256, RpoDigest},
8 Felt, Word, EMPTY_WORD,
9};
10
11mod full;
12pub use full::{Smt, SmtLeaf, SmtLeafError, SmtProof, SmtProofError, SMT_DEPTH};
13
14mod simple;
15pub use simple::SimpleSmt;
16
17// CONSTANTS
18// ================================================================================================
19
20/// Minimum supported depth.
21pub const SMT_MIN_DEPTH: u8 = 1;
22
23/// Maximum supported depth.
24pub const SMT_MAX_DEPTH: u8 = 64;
25
26// SPARSE MERKLE TREE
27// ================================================================================================
28
29/// An abstract description of a sparse Merkle tree.
30///
31/// A sparse Merkle tree is a key-value map which also supports proving that a given value is indeed
32/// stored at a given key in the tree. It is viewed as always being fully populated. If a leaf's
33/// value was not explicitly set, then its value is the default value. Typically, the vast majority
34/// of leaves will store the default value (hence it is "sparse"), and therefore the internal
35/// representation of the tree will only keep track of the leaves that have a different value from
36/// the default.
37///
38/// All leaves sit at the same depth. The deeper the tree, the more leaves it has; but also the
39/// longer its proofs are - of exactly `log(depth)` size. A tree cannot have depth 0, since such a
40/// tree is just a single value, and is probably a programming mistake.
41///
42/// Every key maps to one leaf. If there are as many keys as there are leaves, then
43/// [Self::Leaf] should be the same type as [Self::Value], as is the case with
44/// [crate::merkle::SimpleSmt]. However, if there are more keys than leaves, then [`Self::Leaf`]
45/// must accommodate all keys that map to the same leaf.
46///
47/// [SparseMerkleTree] currently doesn't support optimizations that compress Merkle proofs.
48pub(crate) trait SparseMerkleTree<const DEPTH: u8> {
49 /// The type for a key
50 type Key: Clone + Ord;
51 /// The type for a value
52 type Value: Clone + PartialEq;
53 /// The type for a leaf
54 type Leaf: Clone;
55 /// The type for an opening (i.e. a "proof") of a leaf
56 type Opening;
57
58 /// The default value used to compute the hash of empty leaves
59 const EMPTY_VALUE: Self::Value;
60
61 /// The root of the empty tree with provided DEPTH
62 const EMPTY_ROOT: RpoDigest;
63
64 // PROVIDED METHODS
65 // ---------------------------------------------------------------------------------------------
66
67 /// Returns an opening of the leaf associated with `key`. Conceptually, an opening is a Merkle
68 /// path to the leaf, as well as the leaf itself.
69 fn open(&self, key: &Self::Key) -> Self::Opening {
70 let leaf = self.get_leaf(key);
71
72 let mut index: NodeIndex = {
73 let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(key);
74 leaf_index.into()
75 };
76
77 let merkle_path = {
78 let mut path = Vec::with_capacity(index.depth() as usize);
79 for _ in 0..index.depth() {
80 let is_right = index.is_value_odd();
81 index.move_up();
82 let InnerNode { left, right } = self.get_inner_node(index);
83 let value = if is_right { left } else { right };
84 path.push(value);
85 }
86
87 MerklePath::new(path)
88 };
89
90 Self::path_and_leaf_to_opening(merkle_path, leaf)
91 }
92
93 /// Inserts a value at the specified key, returning the previous value associated with that key.
94 /// Recall that by definition, any key that hasn't been updated is associated with
95 /// [`Self::EMPTY_VALUE`].
96 ///
97 /// This also recomputes all hashes between the leaf (associated with the key) and the root,
98 /// updating the root itself.
99 fn insert(&mut self, key: Self::Key, value: Self::Value) -> Self::Value {
100 let old_value = self.insert_value(key.clone(), value.clone()).unwrap_or(Self::EMPTY_VALUE);
101
102 // if the old value and new value are the same, there is nothing to update
103 if value == old_value {
104 return value;
105 }
106
107 let leaf = self.get_leaf(&key);
108 let node_index = {
109 let leaf_index: LeafIndex<DEPTH> = Self::key_to_leaf_index(&key);
110 leaf_index.into()
111 };
112
113 self.recompute_nodes_from_index_to_root(node_index, Self::hash_leaf(&leaf));
114
115 old_value
116 }
117
118 /// Recomputes the branch nodes (including the root) from `index` all the way to the root.
119 /// `node_hash_at_index` is the hash of the node stored at index.
120 fn recompute_nodes_from_index_to_root(
121 &mut self,
122 mut index: NodeIndex,
123 node_hash_at_index: RpoDigest,
124 ) {
125 let mut node_hash = node_hash_at_index;
126 for node_depth in (0..index.depth()).rev() {
127 let is_right = index.is_value_odd();
128 index.move_up();
129 let InnerNode { left, right } = self.get_inner_node(index);
130 let (left, right) = if is_right {
131 (left, node_hash)
132 } else {
133 (node_hash, right)
134 };
135 node_hash = Rpo256::merge(&[left, right]);
136
137 if node_hash == *EmptySubtreeRoots::entry(DEPTH, node_depth) {
138 // If a subtree is empty, then can remove the inner node, since it's equal to the
139 // default value
140 self.remove_inner_node(index);
141 } else {
142 self.insert_inner_node(index, InnerNode { left, right });
143 }
144 }
145 self.set_root(node_hash);
146 }
147
148 /// Computes what changes are necessary to insert the specified key-value pairs into this Merkle
149 /// tree, allowing for validation before applying those changes.
150 ///
151 /// This method returns a [`MutationSet`], which contains all the information for inserting
152 /// `kv_pairs` into this Merkle tree already calculated, including the new root hash, which can
153 /// be queried with [`MutationSet::root()`]. Once a mutation set is returned,
154 /// [`SparseMerkleTree::apply_mutations()`] can be called in order to commit these changes to
155 /// the Merkle tree, or [`drop()`] to discard them.
156 fn compute_mutations(
157 &self,
158 kv_pairs: impl IntoIterator<Item = (Self::Key, Self::Value)>,
159 ) -> MutationSet<DEPTH, Self::Key, Self::Value> {
160 use NodeMutation::*;
161
162 let mut new_root = self.root();
163 let mut new_pairs: BTreeMap<Self::Key, Self::Value> = Default::default();
164 let mut node_mutations: BTreeMap<NodeIndex, NodeMutation> = Default::default();
165
166 for (key, value) in kv_pairs {
167 // If the old value and the new value are the same, there is nothing to update.
168 // For the unusual case that kv_pairs has multiple values at the same key, we'll have
169 // to check the key-value pairs we've already seen to get the "effective" old value.
170 let old_value = new_pairs.get(&key).cloned().unwrap_or_else(|| self.get_value(&key));
171 if value == old_value {
172 continue;
173 }
174
175 let leaf_index = Self::key_to_leaf_index(&key);
176 let mut node_index = NodeIndex::from(leaf_index);
177
178 // We need the current leaf's hash to calculate the new leaf, but in the rare case that
179 // `kv_pairs` has multiple pairs that go into the same leaf, then those pairs are also
180 // part of the "current leaf".
181 let old_leaf = {
182 let pairs_at_index = new_pairs
183 .iter()
184 .filter(|&(new_key, _)| Self::key_to_leaf_index(new_key) == leaf_index);
185
186 pairs_at_index.fold(self.get_leaf(&key), |acc, (k, v)| {
187 // Most of the time `pairs_at_index` should only contain a single entry (or
188 // none at all), as multi-leaves should be really rare.
189 let existing_leaf = acc.clone();
190 self.construct_prospective_leaf(existing_leaf, k, v)
191 })
192 };
193
194 let new_leaf = self.construct_prospective_leaf(old_leaf, &key, &value);
195
196 let mut new_child_hash = Self::hash_leaf(&new_leaf);
197
198 for node_depth in (0..node_index.depth()).rev() {
199 // Whether the node we're replacing is the right child or the left child.
200 let is_right = node_index.is_value_odd();
201 node_index.move_up();
202
203 let old_node = node_mutations
204 .get(&node_index)
205 .map(|mutation| match mutation {
206 Addition(node) => node.clone(),
207 Removal => EmptySubtreeRoots::get_inner_node(DEPTH, node_depth),
208 })
209 .unwrap_or_else(|| self.get_inner_node(node_index));
210
211 let new_node = if is_right {
212 InnerNode {
213 left: old_node.left,
214 right: new_child_hash,
215 }
216 } else {
217 InnerNode {
218 left: new_child_hash,
219 right: old_node.right,
220 }
221 };
222
223 // The next iteration will operate on this new node's hash.
224 new_child_hash = new_node.hash();
225
226 let &equivalent_empty_hash = EmptySubtreeRoots::entry(DEPTH, node_depth);
227 let is_removal = new_child_hash == equivalent_empty_hash;
228 let new_entry = if is_removal { Removal } else { Addition(new_node) };
229 node_mutations.insert(node_index, new_entry);
230 }
231
232 // Once we're at depth 0, the last node we made is the new root.
233 new_root = new_child_hash;
234 // And then we're done with this pair; on to the next one.
235 new_pairs.insert(key, value);
236 }
237
238 MutationSet {
239 old_root: self.root(),
240 new_root,
241 node_mutations,
242 new_pairs,
243 }
244 }
245
246 /// Applies the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
247 /// this tree.
248 ///
249 /// # Errors
250 /// If `mutations` was computed on a tree with a different root than this one, returns
251 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
252 /// the `mutations` were computed against, and the second item is the actual current root of
253 /// this tree.
254 fn apply_mutations(
255 &mut self,
256 mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
257 ) -> Result<(), MerkleError>
258 where
259 Self: Sized,
260 {
261 use NodeMutation::*;
262 let MutationSet {
263 old_root,
264 node_mutations,
265 new_pairs,
266 new_root,
267 } = mutations;
268
269 // Guard against accidentally trying to apply mutations that were computed against a
270 // different tree, including a stale version of this tree.
271 if old_root != self.root() {
272 return Err(MerkleError::ConflictingRoots {
273 expected_root: self.root(),
274 actual_root: old_root,
275 });
276 }
277
278 for (index, mutation) in node_mutations {
279 match mutation {
280 Removal => {
281 self.remove_inner_node(index);
282 },
283 Addition(node) => {
284 self.insert_inner_node(index, node);
285 },
286 }
287 }
288
289 for (key, value) in new_pairs {
290 self.insert_value(key, value);
291 }
292
293 self.set_root(new_root);
294
295 Ok(())
296 }
297
298 /// Applies the prospective mutations computed with [`SparseMerkleTree::compute_mutations()`] to
299 /// this tree and returns the reverse mutation set. Applying the reverse mutation sets to the
300 /// updated tree will revert the changes.
301 ///
302 /// # Errors
303 /// If `mutations` was computed on a tree with a different root than this one, returns
304 /// [`MerkleError::ConflictingRoots`] with a two-item [`Vec`]. The first item is the root hash
305 /// the `mutations` were computed against, and the second item is the actual current root of
306 /// this tree.
307 fn apply_mutations_with_reversion(
308 &mut self,
309 mutations: MutationSet<DEPTH, Self::Key, Self::Value>,
310 ) -> Result<MutationSet<DEPTH, Self::Key, Self::Value>, MerkleError>
311 where
312 Self: Sized,
313 {
314 use NodeMutation::*;
315 let MutationSet {
316 old_root,
317 node_mutations,
318 new_pairs,
319 new_root,
320 } = mutations;
321
322 // Guard against accidentally trying to apply mutations that were computed against a
323 // different tree, including a stale version of this tree.
324 if old_root != self.root() {
325 return Err(MerkleError::ConflictingRoots {
326 expected_root: self.root(),
327 actual_root: old_root,
328 });
329 }
330
331 let mut reverse_mutations = BTreeMap::new();
332 for (index, mutation) in node_mutations {
333 match mutation {
334 Removal => {
335 if let Some(node) = self.remove_inner_node(index) {
336 reverse_mutations.insert(index, Addition(node));
337 }
338 },
339 Addition(node) => {
340 if let Some(old_node) = self.insert_inner_node(index, node) {
341 reverse_mutations.insert(index, Addition(old_node));
342 } else {
343 reverse_mutations.insert(index, Removal);
344 }
345 },
346 }
347 }
348
349 let mut reverse_pairs = BTreeMap::new();
350 for (key, value) in new_pairs {
351 if let Some(old_value) = self.insert_value(key.clone(), value) {
352 reverse_pairs.insert(key, old_value);
353 } else {
354 reverse_pairs.insert(key, Self::EMPTY_VALUE);
355 }
356 }
357
358 self.set_root(new_root);
359
360 Ok(MutationSet {
361 old_root: new_root,
362 node_mutations: reverse_mutations,
363 new_pairs: reverse_pairs,
364 new_root: old_root,
365 })
366 }
367
368 // REQUIRED METHODS
369 // ---------------------------------------------------------------------------------------------
370
371 /// The root of the tree
372 fn root(&self) -> RpoDigest;
373
374 /// Sets the root of the tree
375 fn set_root(&mut self, root: RpoDigest);
376
377 /// Retrieves an inner node at the given index
378 fn get_inner_node(&self, index: NodeIndex) -> InnerNode;
379
380 /// Inserts an inner node at the given index
381 fn insert_inner_node(&mut self, index: NodeIndex, inner_node: InnerNode) -> Option<InnerNode>;
382
383 /// Removes an inner node at the given index
384 fn remove_inner_node(&mut self, index: NodeIndex) -> Option<InnerNode>;
385
386 /// Inserts a leaf node, and returns the value at the key if already exists
387 fn insert_value(&mut self, key: Self::Key, value: Self::Value) -> Option<Self::Value>;
388
389 /// Returns the value at the specified key. Recall that by definition, any key that hasn't been
390 /// updated is associated with [`Self::EMPTY_VALUE`].
391 fn get_value(&self, key: &Self::Key) -> Self::Value;
392
393 /// Returns the leaf at the specified index.
394 fn get_leaf(&self, key: &Self::Key) -> Self::Leaf;
395
396 /// Returns the hash of a leaf
397 fn hash_leaf(leaf: &Self::Leaf) -> RpoDigest;
398
399 /// Returns what a leaf would look like if a key-value pair were inserted into the tree, without
400 /// mutating the tree itself. The existing leaf can be empty.
401 ///
402 /// To get a prospective leaf based on the current state of the tree, use `self.get_leaf(key)`
403 /// as the argument for `existing_leaf`. The return value from this function can be chained back
404 /// into this function as the first argument to continue making prospective changes.
405 ///
406 /// # Invariants
407 /// Because this method is for a prospective key-value insertion into a specific leaf,
408 /// `existing_leaf` must have the same leaf index as `key` (as determined by
409 /// [`SparseMerkleTree::key_to_leaf_index()`]), or the result will be meaningless.
410 fn construct_prospective_leaf(
411 &self,
412 existing_leaf: Self::Leaf,
413 key: &Self::Key,
414 value: &Self::Value,
415 ) -> Self::Leaf;
416
417 /// Maps a key to a leaf index
418 fn key_to_leaf_index(key: &Self::Key) -> LeafIndex<DEPTH>;
419
420 /// Maps a (MerklePath, Self::Leaf) to an opening.
421 ///
422 /// The length `path` is guaranteed to be equal to `DEPTH`
423 fn path_and_leaf_to_opening(path: MerklePath, leaf: Self::Leaf) -> Self::Opening;
424}
425
426// INNER NODE
427// ================================================================================================
428
429#[derive(Debug, Default, Clone, PartialEq, Eq)]
430#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
431pub struct InnerNode {
432 pub left: RpoDigest,
433 pub right: RpoDigest,
434}
435
436impl InnerNode {
437 pub fn hash(&self) -> RpoDigest {
438 Rpo256::merge(&[self.left, self.right])
439 }
440}
441
442// LEAF INDEX
443// ================================================================================================
444
445/// The index of a leaf, at a depth known at compile-time.
446#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
447#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
448pub struct LeafIndex<const DEPTH: u8> {
449 index: NodeIndex,
450}
451
452impl<const DEPTH: u8> LeafIndex<DEPTH> {
453 pub fn new(value: u64) -> Result<Self, MerkleError> {
454 if DEPTH < SMT_MIN_DEPTH {
455 return Err(MerkleError::DepthTooSmall(DEPTH));
456 }
457
458 Ok(LeafIndex { index: NodeIndex::new(DEPTH, value)? })
459 }
460
461 pub fn value(&self) -> u64 {
462 self.index.value()
463 }
464}
465
466impl LeafIndex<SMT_MAX_DEPTH> {
467 pub const fn new_max_depth(value: u64) -> Self {
468 LeafIndex {
469 index: NodeIndex::new_unchecked(SMT_MAX_DEPTH, value),
470 }
471 }
472}
473
474impl<const DEPTH: u8> From<LeafIndex<DEPTH>> for NodeIndex {
475 fn from(value: LeafIndex<DEPTH>) -> Self {
476 value.index
477 }
478}
479
480impl<const DEPTH: u8> TryFrom<NodeIndex> for LeafIndex<DEPTH> {
481 type Error = MerkleError;
482
483 fn try_from(node_index: NodeIndex) -> Result<Self, Self::Error> {
484 if node_index.depth() != DEPTH {
485 return Err(MerkleError::InvalidNodeIndexDepth {
486 expected: DEPTH,
487 provided: node_index.depth(),
488 });
489 }
490
491 Self::new(node_index.value())
492 }
493}
494
495// MUTATIONS
496// ================================================================================================
497
498/// A change to an inner node of a [`SparseMerkleTree`] that hasn't yet been applied.
499/// [`MutationSet`] stores this type in relation to a [`NodeIndex`] to keep track of what changes
500/// need to occur at which node indices.
501#[derive(Debug, Clone, PartialEq, Eq)]
502pub enum NodeMutation {
503 /// Corresponds to [`SparseMerkleTree::remove_inner_node()`].
504 Removal,
505 /// Corresponds to [`SparseMerkleTree::insert_inner_node()`].
506 Addition(InnerNode),
507}
508
509/// Represents a group of prospective mutations to a `SparseMerkleTree`, created by
510/// `SparseMerkleTree::compute_mutations()`, and that can be applied with
511/// `SparseMerkleTree::apply_mutations()`.
512#[derive(Debug, Clone, PartialEq, Eq, Default)]
513pub struct MutationSet<const DEPTH: u8, K, V> {
514 /// The root of the Merkle tree this MutationSet is for, recorded at the time
515 /// [`SparseMerkleTree::compute_mutations()`] was called. Exists to guard against applying
516 /// mutations to the wrong tree or applying stale mutations to a tree that has since changed.
517 old_root: RpoDigest,
518 /// The set of nodes that need to be removed or added. The "effective" node at an index is the
519 /// Merkle tree's existing node at that index, with the [`NodeMutation`] in this map at that
520 /// index overlayed, if any. Each [`NodeMutation::Addition`] corresponds to a
521 /// [`SparseMerkleTree::insert_inner_node()`] call, and each [`NodeMutation::Removal`]
522 /// corresponds to a [`SparseMerkleTree::remove_inner_node()`] call.
523 node_mutations: BTreeMap<NodeIndex, NodeMutation>,
524 /// The set of top-level key-value pairs we're prospectively adding to the tree, including
525 /// adding empty values. The "effective" value for a key is the value in this BTreeMap, falling
526 /// back to the existing value in the Merkle tree. Each entry corresponds to a
527 /// [`SparseMerkleTree::insert_value()`] call.
528 new_pairs: BTreeMap<K, V>,
529 /// The calculated root for the Merkle tree, given these mutations. Publicly retrievable with
530 /// [`MutationSet::root()`]. Corresponds to a [`SparseMerkleTree::set_root()`]. call.
531 new_root: RpoDigest,
532}
533
534impl<const DEPTH: u8, K, V> MutationSet<DEPTH, K, V> {
535 /// Returns the SMT root that was calculated during `SparseMerkleTree::compute_mutations()`. See
536 /// that method for more information.
537 pub fn root(&self) -> RpoDigest {
538 self.new_root
539 }
540
541 /// Returns the SMT root before the mutations were applied.
542 pub fn old_root(&self) -> RpoDigest {
543 self.old_root
544 }
545
546 /// Returns the set of inner nodes that need to be removed or added.
547 pub fn node_mutations(&self) -> &BTreeMap<NodeIndex, NodeMutation> {
548 &self.node_mutations
549 }
550
551 /// Returns the set of top-level key-value pairs that need to be added, updated or deleted
552 /// (i.e. set to `EMPTY_WORD`).
553 pub fn new_pairs(&self) -> &BTreeMap<K, V> {
554 &self.new_pairs
555 }
556}
557
558// SERIALIZATION
559// ================================================================================================
560
561impl Serializable for InnerNode {
562 fn write_into<W: ByteWriter>(&self, target: &mut W) {
563 self.left.write_into(target);
564 self.right.write_into(target);
565 }
566}
567
568impl Deserializable for InnerNode {
569 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
570 let left = source.read()?;
571 let right = source.read()?;
572
573 Ok(Self { left, right })
574 }
575}
576
577impl Serializable for NodeMutation {
578 fn write_into<W: ByteWriter>(&self, target: &mut W) {
579 match self {
580 NodeMutation::Removal => target.write_bool(false),
581 NodeMutation::Addition(inner_node) => {
582 target.write_bool(true);
583 inner_node.write_into(target);
584 },
585 }
586 }
587}
588
589impl Deserializable for NodeMutation {
590 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
591 if source.read_bool()? {
592 let inner_node = source.read()?;
593 return Ok(NodeMutation::Addition(inner_node));
594 }
595
596 Ok(NodeMutation::Removal)
597 }
598}
599
600impl<const DEPTH: u8, K: Serializable, V: Serializable> Serializable for MutationSet<DEPTH, K, V> {
601 fn write_into<W: ByteWriter>(&self, target: &mut W) {
602 target.write(self.old_root);
603 target.write(self.new_root);
604 self.node_mutations.write_into(target);
605 self.new_pairs.write_into(target);
606 }
607}
608
609impl<const DEPTH: u8, K: Deserializable + Ord, V: Deserializable> Deserializable
610 for MutationSet<DEPTH, K, V>
611{
612 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
613 let old_root = source.read()?;
614 let new_root = source.read()?;
615 let node_mutations = source.read()?;
616 let new_pairs = source.read()?;
617
618 Ok(Self {
619 old_root,
620 node_mutations,
621 new_pairs,
622 new_root,
623 })
624 }
625}