miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.rsRaw
1use super::{
2 mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, KvMap, MerkleError, MerklePath,
3 MerklePathSet, MerkleTree, NodeIndex, RecordingMap, RootPath, Rpo256, RpoDigest, SimpleSmt,
4 TieredSmt, ValuePath, Vec,
5};
6use crate::utils::{
7 collections::{ApplyDiff, Diff, KvMapDiff},
8 ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
9};
10use core::borrow::Borrow;
11
12#[cfg(test)]
13mod tests;
14
15// MERKLE STORE
16// ================================================================================================
17
18/// A default [MerkleStore] which uses a simple [BTreeMap] as the backing storage.
19pub type DefaultMerkleStore = MerkleStore<BTreeMap<RpoDigest, StoreNode>>;
20
21/// A [MerkleStore] with recording capabilities which uses [RecordingMap] as the backing storage.
22pub type RecordingMerkleStore = MerkleStore<RecordingMap<RpoDigest, StoreNode>>;
23
24#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
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)]
93pub struct MerkleStore<T: KvMap<RpoDigest, StoreNode> = BTreeMap<RpoDigest, StoreNode>> {
94 nodes: T,
95}
96
97impl<T: KvMap<RpoDigest, StoreNode>> Default for MerkleStore<T> {
98 fn default() -> Self {
99 Self::new()
100 }
101}
102
103impl<T: KvMap<RpoDigest, StoreNode>> MerkleStore<T> {
104 // CONSTRUCTORS
105 // --------------------------------------------------------------------------------------------
106
107 /// Creates an empty `MerkleStore` instance.
108 pub fn new() -> MerkleStore<T> {
109 // pre-populate the store with the empty hashes
110 let nodes = empty_hashes().into_iter().collect();
111 MerkleStore { nodes }
112 }
113
114 // PUBLIC ACCESSORS
115 // --------------------------------------------------------------------------------------------
116
117 /// Return a count of the non-leaf nodes in the store.
118 pub fn num_internal_nodes(&self) -> usize {
119 self.nodes.len()
120 }
121
122 /// Returns the node at `index` rooted on the tree `root`.
123 ///
124 /// # Errors
125 /// This method can return the following errors:
126 /// - `RootNotInStore` if the `root` is not present in the store.
127 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
128 /// the store.
129 pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
130 let mut hash = root;
131
132 // corner case: check the root is in the store when called with index `NodeIndex::root()`
133 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
134
135 for i in (0..index.depth()).rev() {
136 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
137
138 let bit = (index.value() >> i) & 1;
139 hash = if bit == 0 { node.left } else { node.right }
140 }
141
142 Ok(hash)
143 }
144
145 /// Returns the node at the specified `index` and its opening to the `root`.
146 ///
147 /// The path starts at the sibling of the target leaf.
148 ///
149 /// # Errors
150 /// This method can return the following errors:
151 /// - `RootNotInStore` if the `root` is not present in the store.
152 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
153 /// the store.
154 pub fn get_path(&self, root: RpoDigest, index: NodeIndex) -> Result<ValuePath, MerkleError> {
155 let mut hash = root;
156 let mut path = Vec::with_capacity(index.depth().into());
157
158 // corner case: check the root is in the store when called with index `NodeIndex::root()`
159 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
160
161 for i in (0..index.depth()).rev() {
162 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
163
164 let bit = (index.value() >> i) & 1;
165 hash = if bit == 0 {
166 path.push(node.right);
167 node.left
168 } else {
169 path.push(node.left);
170 node.right
171 }
172 }
173
174 // the path is computed from root to leaf, so it must be reversed
175 path.reverse();
176
177 Ok(ValuePath {
178 value: hash,
179 path: MerklePath::new(path),
180 })
181 }
182
183 /// Reconstructs a path from the root until a leaf or empty node and returns its depth.
184 ///
185 /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
186 /// from `root`. The maximum value the argument accepts is [u64::BITS].
187 ///
188 /// The traversed path from leaf to root will start at the least significant bit of `index`,
189 /// and will be executed for `tree_depth` bits.
190 ///
191 /// # Errors
192 /// Will return an error if:
193 /// - The provided root is not found.
194 /// - The path from the root continues to a depth greater than `tree_depth`.
195 /// - The provided `tree_depth` is greater than `64.
196 /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
197 /// information, check [NodeIndex::new].
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 // it's not illegal to have a maximum depth of `0`; we should just return the root in that
211 // case. this check will simplify the implementation as we could overflow bits for depth
212 // `0`.
213 if tree_depth == 0 {
214 return Ok(0);
215 }
216
217 // check if the root exists, providing the proper error report if it doesn't
218 let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
219 let mut hash = root;
220 if !self.nodes.contains_key(&hash) {
221 return Err(MerkleError::RootNotInStore(hash));
222 }
223
224 // we traverse from root to leaf, so the path is reversed
225 let mut path = (index << (64 - tree_depth)).reverse_bits();
226
227 // iterate every depth and reconstruct the path from root to leaf
228 for depth in 0..tree_depth {
229 // we short-circuit if an empty node has been found
230 if hash == empty[depth as usize] {
231 return Ok(depth);
232 }
233
234 // fetch the children pair, mapped by its parent hash
235 let children = match self.nodes.get(&hash) {
236 Some(node) => node,
237 None => return Ok(depth),
238 };
239
240 // traverse down
241 hash = if path & 1 == 0 { children.left } else { children.right };
242 path >>= 1;
243 }
244
245 // at max depth assert it doesn't have sub-trees
246 if self.nodes.contains_key(&hash) {
247 return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
248 }
249
250 // depleted bits; return max depth
251 Ok(tree_depth)
252 }
253
254 // DATA EXTRACTORS
255 // --------------------------------------------------------------------------------------------
256
257 /// Returns a subset of this Merkle store such that the returned Merkle store contains all
258 /// nodes which are descendants of the specified roots.
259 ///
260 /// The roots for which no descendants exist in this Merkle store are ignored.
261 pub fn subset<I, R>(&self, roots: I) -> MerkleStore<T>
262 where
263 I: Iterator<Item = R>,
264 R: Borrow<RpoDigest>,
265 {
266 let mut store = MerkleStore::new();
267 for root in roots {
268 let root = *root.borrow();
269 store.clone_tree_from(root, self);
270 }
271 store
272 }
273
274 /// Iterator over the inner nodes of the [MerkleStore].
275 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
276 self.nodes.iter().map(|(r, n)| InnerNodeInfo {
277 value: *r,
278 left: n.left,
279 right: n.right,
280 })
281 }
282
283 // STATE MUTATORS
284 // --------------------------------------------------------------------------------------------
285
286 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
287 /// new root.
288 ///
289 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
290 /// include all the nodes into the store.
291 pub fn add_merkle_path(
292 &mut self,
293 index: u64,
294 node: RpoDigest,
295 path: MerklePath,
296 ) -> Result<RpoDigest, MerkleError> {
297 let root = path.inner_nodes(index, node)?.fold(RpoDigest::default(), |_, node| {
298 let value: RpoDigest = node.value;
299 let left: RpoDigest = node.left;
300 let right: RpoDigest = node.right;
301
302 debug_assert_eq!(Rpo256::merge(&[left, right]), value);
303 self.nodes.insert(value, StoreNode { left, right });
304
305 node.value
306 });
307 Ok(root)
308 }
309
310 /// Adds all the nodes of multiple Merkle paths into the store.
311 ///
312 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
313 /// into the store.
314 ///
315 /// For further reference, check [MerkleStore::add_merkle_path].
316 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<(), MerkleError>
317 where
318 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
319 {
320 for (index_value, node, path) in paths.into_iter() {
321 self.add_merkle_path(index_value, node, path)?;
322 }
323 Ok(())
324 }
325
326 /// Appends the provided [MerklePathSet] into the store.
327 ///
328 /// For further reference, check [MerkleStore::add_merkle_path].
329 pub fn add_merkle_path_set(
330 &mut self,
331 path_set: &MerklePathSet,
332 ) -> Result<RpoDigest, MerkleError> {
333 let root = path_set.root();
334 for (index, path) in path_set.to_paths() {
335 self.add_merkle_path(index, path.value, path.path)?;
336 }
337 Ok(root)
338 }
339
340 /// Sets a node to `value`.
341 ///
342 /// # Errors
343 /// This method can return the following errors:
344 /// - `RootNotInStore` if the `root` is not present in the store.
345 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in
346 /// the store.
347 pub fn set_node(
348 &mut self,
349 mut root: RpoDigest,
350 index: NodeIndex,
351 value: RpoDigest,
352 ) -> Result<RootPath, MerkleError> {
353 let node = value;
354 let ValuePath { value, path } = self.get_path(root, index)?;
355
356 // performs the update only if the node value differs from the opening
357 if node != value {
358 root = self.add_merkle_path(index.value(), node, path.clone())?;
359 }
360
361 Ok(RootPath { root, path })
362 }
363
364 /// Merges two elements and adds the resulting node into the store.
365 ///
366 /// Merges arbitrary values. They may be leafs, nodes, or a mixture of both.
367 pub fn merge_roots(
368 &mut self,
369 left_root: RpoDigest,
370 right_root: RpoDigest,
371 ) -> Result<RpoDigest, MerkleError> {
372 let parent = Rpo256::merge(&[left_root, right_root]);
373 self.nodes.insert(
374 parent,
375 StoreNode {
376 left: left_root,
377 right: right_root,
378 },
379 );
380
381 Ok(parent)
382 }
383
384 // DESTRUCTURING
385 // --------------------------------------------------------------------------------------------
386
387 /// Returns the inner storage of this MerkleStore while consuming `self`.
388 pub fn into_inner(self) -> T {
389 self.nodes
390 }
391
392 // HELPER METHODS
393 // --------------------------------------------------------------------------------------------
394
395 /// Recursively clones a tree with the specified root from the specified source into self.
396 ///
397 /// If the source store does not contain a tree with the specified root, this is a noop.
398 fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) {
399 // process the node only if it is in the source
400 if let Some(node) = source.nodes.get(&root) {
401 // if the node has already been inserted, no need to process it further as all of its
402 // descendants should be already cloned from the source store
403 if self.nodes.insert(root, *node).is_none() {
404 self.clone_tree_from(node.left, source);
405 self.clone_tree_from(node.right, source);
406 }
407 }
408 }
409}
410
411// CONVERSIONS
412// ================================================================================================
413
414impl<T: KvMap<RpoDigest, StoreNode>> From<&MerkleTree> for MerkleStore<T> {
415 fn from(value: &MerkleTree) -> Self {
416 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
417 Self { nodes }
418 }
419}
420
421impl<T: KvMap<RpoDigest, StoreNode>> From<&SimpleSmt> for MerkleStore<T> {
422 fn from(value: &SimpleSmt) -> Self {
423 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
424 Self { nodes }
425 }
426}
427
428impl<T: KvMap<RpoDigest, StoreNode>> From<&Mmr> for MerkleStore<T> {
429 fn from(value: &Mmr) -> Self {
430 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
431 Self { nodes }
432 }
433}
434
435impl<T: KvMap<RpoDigest, StoreNode>> From<&TieredSmt> for MerkleStore<T> {
436 fn from(value: &TieredSmt) -> Self {
437 let nodes = combine_nodes_with_empty_hashes(value.inner_nodes()).collect();
438 Self { nodes }
439 }
440}
441
442impl<T: KvMap<RpoDigest, StoreNode>> From<T> for MerkleStore<T> {
443 fn from(values: T) -> Self {
444 let nodes = values.into_iter().chain(empty_hashes()).collect();
445 Self { nodes }
446 }
447}
448
449impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<InnerNodeInfo> for MerkleStore<T> {
450 fn from_iter<I: IntoIterator<Item = InnerNodeInfo>>(iter: I) -> Self {
451 let nodes = combine_nodes_with_empty_hashes(iter).collect();
452 Self { nodes }
453 }
454}
455
456impl<T: KvMap<RpoDigest, StoreNode>> FromIterator<(RpoDigest, StoreNode)> for MerkleStore<T> {
457 fn from_iter<I: IntoIterator<Item = (RpoDigest, StoreNode)>>(iter: I) -> Self {
458 let nodes = iter.into_iter().chain(empty_hashes()).collect();
459 Self { nodes }
460 }
461}
462
463// ITERATORS
464// ================================================================================================
465
466impl<T: KvMap<RpoDigest, StoreNode>> Extend<InnerNodeInfo> for MerkleStore<T> {
467 fn extend<I: IntoIterator<Item = InnerNodeInfo>>(&mut self, iter: I) {
468 self.nodes.extend(iter.into_iter().map(|info| {
469 (
470 info.value,
471 StoreNode {
472 left: info.left,
473 right: info.right,
474 },
475 )
476 }));
477 }
478}
479
480// DiffT & ApplyDiffT TRAIT IMPLEMENTATION
481// ================================================================================================
482impl<T: KvMap<RpoDigest, StoreNode>> Diff<RpoDigest, StoreNode> for MerkleStore<T> {
483 type DiffType = KvMapDiff<RpoDigest, StoreNode>;
484
485 fn diff(&self, other: &Self) -> Self::DiffType {
486 self.nodes.diff(&other.nodes)
487 }
488}
489
490impl<T: KvMap<RpoDigest, StoreNode>> ApplyDiff<RpoDigest, StoreNode> for MerkleStore<T> {
491 type DiffType = KvMapDiff<RpoDigest, StoreNode>;
492
493 fn apply(&mut self, diff: Self::DiffType) {
494 self.nodes.apply(diff);
495 }
496}
497
498// SERIALIZATION
499// ================================================================================================
500
501impl Serializable for StoreNode {
502 fn write_into<W: ByteWriter>(&self, target: &mut W) {
503 self.left.write_into(target);
504 self.right.write_into(target);
505 }
506}
507
508impl Deserializable for StoreNode {
509 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
510 let left = RpoDigest::read_from(source)?;
511 let right = RpoDigest::read_from(source)?;
512 Ok(StoreNode { left, right })
513 }
514}
515
516impl<T: KvMap<RpoDigest, StoreNode>> Serializable for MerkleStore<T> {
517 fn write_into<W: ByteWriter>(&self, target: &mut W) {
518 target.write_u64(self.nodes.len() as u64);
519
520 for (k, v) in self.nodes.iter() {
521 k.write_into(target);
522 v.write_into(target);
523 }
524 }
525}
526
527impl<T: KvMap<RpoDigest, StoreNode>> Deserializable for MerkleStore<T> {
528 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
529 let len = source.read_u64()?;
530 let mut nodes: Vec<(RpoDigest, StoreNode)> = Vec::with_capacity(len as usize);
531
532 for _ in 0..len {
533 let key = RpoDigest::read_from(source)?;
534 let value = StoreNode::read_from(source)?;
535 nodes.push((key, value));
536 }
537
538 Ok(nodes.into_iter().collect())
539 }
540}
541
542// HELPER FUNCTIONS
543// ================================================================================================
544
545/// Creates empty hashes for all the subtrees of a tree with a max depth of 255.
546fn empty_hashes() -> impl IntoIterator<Item = (RpoDigest, StoreNode)> {
547 let subtrees = EmptySubtreeRoots::empty_hashes(255);
548 subtrees.iter().rev().copied().zip(subtrees.iter().rev().skip(1).copied()).map(
549 |(child, parent)| {
550 (
551 parent,
552 StoreNode {
553 left: child,
554 right: child,
555 },
556 )
557 },
558 )
559}
560
561/// Consumes an iterator of [InnerNodeInfo] and returns an iterator of `(value, node)` tuples
562/// which includes the nodes associate with roots of empty subtrees up to a depth of 255.
563fn combine_nodes_with_empty_hashes(
564 nodes: impl IntoIterator<Item = InnerNodeInfo>,
565) -> impl Iterator<Item = (RpoDigest, StoreNode)> {
566 nodes
567 .into_iter()
568 .map(|info| {
569 (
570 info.value,
571 StoreNode {
572 left: info.left,
573 right: info.right,
574 },
575 )
576 })
577 .chain(empty_hashes())
578}