miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.rsRaw
1use alloc::{collections::BTreeMap, vec::Vec};
2use core::borrow::Borrow;
3
4use super::{
5 mmr::Mmr, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerkleTree, NodeIndex,
6 PartialMerkleTree, RootPath, Rpo256, RpoDigest, SimpleSmt, Smt, ValuePath,
7};
8use crate::utils::{
9 collections::{KvMap, RecordingMap},
10 ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
11};
12
13#[cfg(test)]
14mod tests;
15
16// MERKLE STORE
17// ================================================================================================
18
19/// A default [MerkleStore] which uses a simple [BTreeMap] as the backing storage.
20pub type DefaultMerkleStore = MerkleStore<BTreeMap<RpoDigest, StoreNode>>;
21
22/// A [MerkleStore] with recording capabilities which uses [RecordingMap] as the backing storage.
23pub type RecordingMerkleStore = MerkleStore<RecordingMap<RpoDigest, StoreNode>>;
24
25#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
26#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
27pub struct StoreNode {
28 left: RpoDigest,
29 right: RpoDigest,
30}
31
32/// An in-memory data store for Merkelized data.
33///
34/// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple
35/// trees to live as long as necessary and without duplication, this allows the implementation of
36/// space efficient persistent data structures.
37///
38/// Example usage:
39///
40/// ```rust
41/// # use miden_crypto::{ZERO, Felt, Word};
42/// # use miden_crypto::merkle::{NodeIndex, MerkleStore, MerkleTree};
43/// # use miden_crypto::hash::rpo::Rpo256;
44/// # const fn int_to_node(value: u64) -> Word {
45/// # [Felt::new(value), ZERO, ZERO, ZERO]
46/// # }
47/// # let A = int_to_node(1);
48/// # let B = int_to_node(2);
49/// # let C = int_to_node(3);
50/// # let D = int_to_node(4);
51/// # let E = int_to_node(5);
52/// # let F = int_to_node(6);
53/// # let G = int_to_node(7);
54/// # let H0 = int_to_node(8);
55/// # let H1 = int_to_node(9);
56/// # let T0 = MerkleTree::new([A, B, C, D, E, F, G, H0].to_vec()).expect("even number of leaves provided");
57/// # let T1 = MerkleTree::new([A, B, C, D, E, F, G, H1].to_vec()).expect("even number of leaves provided");
58/// # let ROOT0 = T0.root();
59/// # let ROOT1 = T1.root();
60/// let mut store: MerkleStore = MerkleStore::new();
61///
62/// // the store is initialized with the SMT empty nodes
63/// assert_eq!(store.num_internal_nodes(), 255);
64///
65/// let tree1 = MerkleTree::new(vec![A, B, C, D, E, F, G, H0]).unwrap();
66/// let tree2 = MerkleTree::new(vec![A, B, C, D, E, F, G, H1]).unwrap();
67///
68/// // populates the store with two merkle trees, common nodes are shared
69/// store.extend(tree1.inner_nodes());
70/// store.extend(tree2.inner_nodes());
71///
72/// // every leaf except the last are the same
73/// for i in 0..7 {
74/// let idx0 = NodeIndex::new(3, i).unwrap();
75/// let d0 = store.get_node(ROOT0, idx0).unwrap();
76/// let idx1 = NodeIndex::new(3, i).unwrap();
77/// let d1 = store.get_node(ROOT1, idx1).unwrap();
78/// assert_eq!(d0, d1, "Both trees have the same leaf at pos {i}");
79/// }
80///
81/// // The leafs A-B-C-D are the same for both trees, so are their 2 immediate parents
82/// for i in 0..4 {
83/// let idx0 = NodeIndex::new(3, i).unwrap();
84/// let d0 = store.get_path(ROOT0, idx0).unwrap();
85/// let idx1 = NodeIndex::new(3, i).unwrap();
86/// let d1 = store.get_path(ROOT1, idx1).unwrap();
87/// assert_eq!(d0.path[0..2], d1.path[0..2], "Both sub-trees are equal up to two levels");
88/// }
89///
90/// // Common internal nodes are shared, the two added trees have a total of 30, but the store has
91/// // only 10 new entries, corresponding to the 10 unique internal nodes of these trees.
92/// assert_eq!(store.num_internal_nodes() - 255, 10);
93/// ```
94#[derive(Debug, Clone, Eq, PartialEq)]
95#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
96pub struct MerkleStore<T: KvMap<RpoDigest, StoreNode> = BTreeMap<RpoDigest, StoreNode>> {
97 nodes: T,
98}
99
100impl<T: KvMap<RpoDigest, StoreNode>> Default for MerkleStore<T> {
101 fn default() -> Self {
102 Self::new()
103 }
104}
105
106impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
107 // CONSTRUCTORS
108 // --------------------------------------------------------------------------------------------
109
110 /// Creates an empty `MerkleStore` instance.
111 pub fn new() -> MerkleStore<T> {
112 // pre-populate the store with the empty hashes
113 let nodes = empty_hashes().into_iter().collect();
114 MerkleStore { nodes }
115 }
116
117 // PUBLIC ACCESSORS
118 // --------------------------------------------------------------------------------------------
119
120 /// Return a count of the non-leaf nodes in the store.
121 pub fn num_internal_nodes(&self) -> usize {
122 self.nodes.len()
123 }
124
125 /// Returns the node at `index` rooted on the tree `root`.
126 ///
127 /// # Errors
128 /// This method can return the following errors:
129 /// - `RootNotInStore` if the `root` is not present in the store.
130 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
131 /// the store.
132 pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
133 let mut hash = root;
134
135 // corner case: check the root is in the store when called with index `NodeIndex::root()`
136 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
137
138 for i in (0..index.depth()).rev() {
139 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
140
141 let bit = (index.value() >> i) & 1;
142 hash = if bit == 0 { node.left } else { node.right }
143 }
144
145 Ok(hash)
146 }
147
148 /// Returns the node at the specified `index` and its opening to the `root`.
149 ///
150 /// The path starts at the sibling of the target leaf.
151 ///
152 /// # Errors
153 /// This method can return the following errors:
154 /// - `RootNotInStore` if the `root` is not present in the store.
155 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
156 /// the store.
157 pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result<ValuePath, MerkleError> {
158 let mut hash = root;
159 let mut path = Vec::with_capacity(index.depth().into());
160
161 // corner case: check the root is in the store when called with index `NodeIndex::root()`
162 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
163
164 for i in (0..index.depth()).rev() {
165 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
166
167 let bit = (index.value() >> i) & 1;
168 hash = if bit == 0 {
169 path.push(node.right);
170 node.left
171 } else {
172 path.push(node.left);
173 node.right
174 }
175 }
176
177 // the path is computed from root to leaf, so it must be reversed
178 path.reverse();
179
180 Ok(ValuePath::new(hash, MerklePath::new(path)))
181 }
182
183 // LEAF TRAVERSAL
184 // --------------------------------------------------------------------------------------------
185
186 /// Returns the depth of the first leaf or an empty node encountered while traversing the tree
187 /// from the specified root down according to the provided index.
188 ///
189 /// The `tree_depth` parameter specifies the depth of the tree rooted at `root`. The
190 /// maximum value the argument accepts is [u64::BITS].
191 ///
192 /// # Errors
193 /// Will return an error if:
194 /// - The provided root is not found.
195 /// - The provided `tree_depth` is greater than 64.
196 /// - The provided `index` is not valid for a depth equivalent to `tree_depth`.
197 /// - No leaf or an empty node was found while traversing the tree down to `tree_depth`.
198 pub fn get_leaf_depth(
199 &self,
200 root: RpoDigest,
201 tree_depth: u8,
202 index: u64,
203 ) -> Result<u8, MerkleError> {
204 // validate depth and index
205 if tree_depth > 64 {
206 return Err(MerkleError::DepthTooBig(tree_depth as u64));
207 }
208 NodeIndex::new(tree_depth, index)?;
209
210 // check if the root exists, providing the proper error report if it doesn't
211 let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
212 let mut hash = root;
213 if !self.nodes.contains_key(&hash) {
214 return Err(MerkleError::RootNotInStore(hash));
215 }
216
217 // we traverse from root to leaf, so the path is reversed
218 let mut path = (index << (64 - tree_depth)).reverse_bits();
219
220 // iterate every depth and reconstruct the path from root to leaf
221 for depth in 0..=tree_depth {
222 // we short-circuit if an empty node has been found
223 if hash == empty[depth as usize] {
224 return Ok(depth);
225 }
226
227 // fetch the children pair, mapped by its parent hash
228 let children = match self.nodes.get(&hash) {
229 Some(node) => node,
230 None => return Ok(depth),
231 };
232
233 // traverse down
234 hash = if path & 1 == 0 { children.left } else { children.right };
235 path >>= 1;
236 }
237
238 // return an error because we exhausted the index but didn't find either a leaf or an
239 // empty node
240 Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
241 }
242
243 /// Returns index and value of a leaf node which is the only leaf node in a subtree defined by
244 /// the provided root. If the subtree contains zero or more than one leaf nodes None is
245 /// returned.
246 ///
247 /// The `tree_depth` parameter specifies the depth of the parent tree such that `root` is
248 /// located in this tree at `root_index`. The maximum value the argument accepts is
249 /// [u64::BITS].
250 ///
251 /// # Errors
252 /// Will return an error if:
253 /// - The provided root is not found.
254 /// - The provided `tree_depth` is greater than 64.
255 /// - The provided `root_index` has depth greater than `tree_depth`.
256 /// - A lone node at depth `tree_depth` is not a leaf node.
257 pub fn find_lone_leaf(
258 &self,
259 root: RpoDigest,
260 root_index: NodeIndex,
261 tree_depth: u8,
262 ) -> Result<Option<(NodeIndex, RpoDigest)>, MerkleError> {
263 // we set max depth at u64::BITS as this is the largest meaningful value for a 64-bit index
264 const MAX_DEPTH: u8 = u64::BITS as u8;
265 if tree_depth > MAX_DEPTH {
266 return Err(MerkleError::DepthTooBig(tree_depth as u64));
267 }
268 let empty = EmptySubtreeRoots::empty_hashes(MAX_DEPTH);
269
270 let mut node = root;
271 if !self.nodes.contains_key(&node) {
272 return Err(MerkleError::RootNotInStore(node));
273 }
274
275 let mut index = root_index;
276 if index.depth() > tree_depth {
277 return Err(MerkleError::DepthTooBig(index.depth() as u64));
278 }
279
280 // traverse down following the path of single non-empty nodes; this works because if a
281 // node has two empty children it cannot contain a lone leaf. similarly if a node has
282 // two non-empty children it must contain at least two leaves.
283 for depth in index.depth()..tree_depth {
284 // if the node is a leaf, return; otherwise, examine the node's children
285 let children = match self.nodes.get(&node) {
286 Some(node) => node,
287 None => return Ok(Some((index, node))),
288 };
289
290 let empty_node = empty[depth as usize + 1];
291 node = if children.left != empty_node && children.right == empty_node {
292 index = index.left_child();
293 children.left
294 } else if children.left == empty_node && children.right != empty_node {
295 index = index.right_child();
296 children.right
297 } else {
298 return Ok(None);
299 };
300 }
301
302 // if we are here, we got to `tree_depth`; thus, either the current node is a leaf node,
303 // and so we return it, or it is an internal node, and then we return an error
304 if self.nodes.contains_key(&node) {
305 Err(MerkleError::DepthTooBig(tree_depth as u64 + 1))
306 } else {
307 Ok(Some((index, node)))
308 }
309 }
310
311 // DATA EXTRACTORS
312 // --------------------------------------------------------------------------------------------
313
314 /// Returns a subset of this Merkle store such that the returned Merkle store contains all
315 /// nodes which are descendants of the specified roots.
316 ///
317 /// The roots for which no descendants exist in this Merkle store are ignored.
318 pub fn subset<I, R>(&self, roots: I) -> MerkleStore<T>
319 where
320 I: Iterator<Item = R>,
321 R: Borrow<RpoDigest>,
322 {
323 let mut store = MerkleStore::new();
324 for root in roots {
325 let root = *root.borrow();
326 store.clone_tree_from(root, self);
327 }
328 store
329 }
330
331 /// Iterator over the inner nodes of the [MerkleStore].
332 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
333 self.nodes
334 .iter()
335 .map(|(r, n)| InnerNodeInfo { value: *r, left: n.left, right: n.right })
336 }
337
338 /// Iterator over the non-empty leaves of the Merkle tree associated with the specified `root`
339 /// and `max_depth`.
340 pub fn non_empty_leaves(
341 &self,
342 root: RpoDigest,
343 max_depth: u8,
344 ) -> impl Iterator<Item = (NodeIndex, RpoDigest)> + '_ {
345 let empty_roots = EmptySubtreeRoots::empty_hashes(max_depth);
346 let mut stack = Vec::new();
347 stack.push((NodeIndex::new_unchecked(0, 0), root));
348
349 core::iter::from_fn(move || {
350 while let Some((index, node_hash)) = stack.pop() {
351 // if we are at the max depth then we have reached a leaf
352 if index.depth() == max_depth {
353 return Some((index, node_hash));
354 }
355
356 // fetch the nodes children and push them onto the stack if they are not the roots
357 // of empty subtrees
358 if let Some(node) = self.nodes.get(&node_hash) {
359 if !empty_roots.contains(&node.left) {
360 stack.push((index.left_child(), node.left));
361 }
362 if !empty_roots.contains(&node.right) {
363 stack.push((index.right_child(), node.right));
364 }
365
366 // if the node is not in the store assume it is a leaf
367 } else {
368 return Some((index, node_hash));
369 }
370 }
371
372 None
373 })
374 }
375
376 // STATE MUTATORS
377 // --------------------------------------------------------------------------------------------
378
379 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
380 /// new root.
381 ///
382 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
383 /// include all the nodes into the store.
384 pub fn add_merkle_path(
385 &mut self,
386 index: u64,
387 node: RpoDigest,
388 path: MerklePath,
389 ) -> Result<RpoDigest, MerkleError> {
390 let root = path.inner_nodes(index, node)?.fold(RpoDigest::default(), |_, node| {
391 let value: RpoDigest = node.value;
392 let left: RpoDigest = node.left;
393 let right: RpoDigest = node.right;
394
395 debug_assert_eq!(Rpo256::merge(&[left, right]), value);
396 self.nodes.insert(value, StoreNode { left, right });
397
398 node.value
399 });
400 Ok(root)
401 }
402
403 /// Adds all the nodes of multiple Merkle paths into the store.
404 ///
405 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
406 /// into the store.
407 ///
408 /// For further reference, check [MerkleStore::add_merkle_path].
409 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<(), MerkleError>
410 where
411 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
412 {
413 for (index_value, node, path) in paths.into_iter() {
414 self.add_merkle_path(index_value, node, path)?;
415 }
416 Ok(())
417 }
418
419 /// Sets a node to `value`.
420 ///
421 /// # Errors
422 /// This method can return the following errors:
423 /// - `RootNotInStore` if the `root` is not present in the store.
424 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
425 /// the store.
426 pub fn set_node(
427 &mut self,
428 mut root: RpoDigest,
429 index: NodeIndex,
430 value: RpoDigest,
431 ) -> Result<RootPath, MerkleError> {
432 let node = value;
433 let ValuePath { value, path } = self.get_path(root, index)?;
434
435 // performs the update only if the node value differs from the opening
436 if node != value {
437 root = self.add_merkle_path(index.value(), node, path.clone())?;
438 }
439
440 Ok(RootPath { root, path })
441 }
442
443 /// Merges two elements and adds the resulting node into the store.
444 ///
445 /// Merges arbitrary values. They may be leafs, nodes, or a mixture of both.
446 pub fn merge_roots(
447 &mut self,
448 left_root: RpoDigest,
449 right_root: RpoDigest,
450 ) -> Result<RpoDigest, MerkleError> {
451 let parent = Rpo256::merge(&[left_root, right_root]);
452 self.nodes.insert(parent, StoreNode { left: left_root, right: right_root });
453
454 Ok(parent)
455 }
456
457 // DESTRUCTURING
458 // --------------------------------------------------------------------------------------------
459
460 /// Returns the inner storage of this MerkleStore while consuming `self`.
461 pub fn into_inner(self) -> T {
462 self.nodes
463 }
464
465 // HELPER METHODS
466 // --------------------------------------------------------------------------------------------
467
468 /// Recursively clones a tree with the specified root from the specified source into self.
469 ///
470 /// If the source store does not contain a tree with the specified root, this is a noop.
471 fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) {
472 // process the node only if it is in the source
473 if let Some(node) = source.nodes.get(&root) {
474 // if the node has already been inserted, no need to process it further as all of its
475 // descendants should be already cloned from the source store
476 if self.nodes.insert(root, *node).is_none() {
477 self.clone_tree_from(node.left, source);
478 self.clone_tree_from(node.right, source);
479 }
480 }
481 }
482}
483
484// CONVERSIONS
485// ================================================================================================
486
487impl<T: KvMap<RpoDigest, StoreNode>> From<&MerkleTree> for MerkleStore<T> {
488 fn from(value: &MerkleTree) -> Self {
489 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
490 Self { nodes }
491 }
492}
493
494impl<T: KvMap<RpoDigest, StoreNode>, const DEPTH: u8> From<&SimpleSmt<DEPTH>> for MerkleStore<T> {
495 fn from(value: &SimpleSmt<DEPTH>) -> Self {
496 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
497 Self { nodes }
498 }
499}
500
501impl<T: KvMap<RpoDigest, StoreNode>> From<&Smt> for MerkleStore<T> {
502 fn from(value: &Smt) -> Self {
503 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
504 Self { nodes }
505 }
506}
507
508impl<T: KvMap<RpoDigest, StoreNode>> From<&Mmr> for MerkleStore<T> {
509 fn from(value: &Mmr) -> Self {
510 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
511 Self { nodes }
512 }
513}
514
515impl<T: KvMap<RpoDigest, StoreNode>> From<&PartialMerkleTree> for MerkleStore<T> {
516 fn from(value: &PartialMerkleTree) -> Self {
517 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
518 Self { nodes }
519 }
520}
521
522impl<T: KvMap<RpoDigest, StoreNode>> From<T> for MerkleStore<T> {
523 fn from(values: T) -> Self {
524 let nodes = values.into_iter().chain(empty_hashes()).collect();
525 Self { nodes }
526 }
527}
528
529impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<InnerNodeInfo> for MerkleStore<T> {
530 fn from_iter<I: IntoIterator<Item = InnerNodeInfo>>(iter: I) -> Self {
531 let nodes = combine_nodes_with_empty_hashes(iter).collect();
532 Self { nodes }
533 }
534}
535
536impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<(RpoDigest, StoreNode)> for MerkleStore<T> {
537 fn from_iter<I: IntoIterator<Item = (RpoDigest, StoreNode)>>(iter: I) -> Self {
538 let nodes = iter.into_iter().chain(empty_hashes()).collect();
539 Self { nodes }
540 }
541}
542
543// ITERATORS
544// ================================================================================================
545impl<T: KvMap<RpoDigest, StoreNode>> Extend<InnerNodeInfo> for MerkleStore<T> {
546 fn extend<I: IntoIterator<Item = InnerNodeInfo>>(&mut self, iter: I) {
547 self.nodes.extend(
548 iter.into_iter()
549 .map(|info| (info.value, StoreNode { left: info.left, right: info.right })),
550 );
551 }
552}
553
554// SERIALIZATION
555// ================================================================================================
556
557impl Serializable for StoreNode {
558 fn write_into<W: ByteWriter>(&self, target: &mut W) {
559 self.left.write_into(target);
560 self.right.write_into(target);
561 }
562}
563
564impl Deserializable for StoreNode {
565 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
566 let left = RpoDigest::read_from(source)?;
567 let right = RpoDigest::read_from(source)?;
568 Ok(StoreNode { left, right })
569 }
570}
571
572impl<T: KvMap<RpoDigest, StoreNode>> Serializable for MerkleStore<T> {
573 fn write_into<W: ByteWriter>(&self, target: &mut W) {
574 target.write_u64(self.nodes.len() as u64);
575
576 for (k, v) in self.nodes.iter() {
577 k.write_into(target);
578 v.write_into(target);
579 }
580 }
581}
582
583impl<T: KvMap<RpoDigest, StoreNode>> Deserializable for MerkleStore<T> {
584 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
585 let len = source.read_u64()?;
586 let mut nodes: Vec<(RpoDigest, StoreNode)> = Vec::with_capacity(len as usize);
587
588 for _ in 0..len {
589 let key = RpoDigest::read_from(source)?;
590 let value = StoreNode::read_from(source)?;
591 nodes.push((key, value));
592 }
593
594 Ok(nodes.into_iter().collect())
595 }
596}
597
598// HELPER FUNCTIONS
599// ================================================================================================
600
601/// Creates empty hashes for all the subtrees of a tree with a max depth of 255.
602fn empty_hashes() -> impl IntoIterator<Item = (RpoDigest, StoreNode)> {
603 let subtrees = EmptySubtreeRoots::empty_hashes(255);
604 subtrees
605 .iter()
606 .rev()
607 .copied()
608 .zip(subtrees.iter().rev().skip(1).copied())
609 .map(|(child, parent)| (parent, StoreNode { left: child, right: child }))
610}
611
612/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples
613/// which includes the nodes associate with roots of empty subtrees up to a depth of 255.
614fn combine_nodes_with_empty_hashes(
615 nodes: impl IntoIterator<Item = InnerNodeInfo>,
616) -> impl Iterator<Item = (RpoDigest, StoreNode)> {
617 nodes
618 .into_iter()
619 .map(|info| (info.value, StoreNode { left: info.left, right: info.right }))
620 .chain(empty_hashes())
621}