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