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