miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.rsRaw
1use super::{
2 mmr::Mmr, BTreeMap, EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, MerklePathSet,
3 MerkleTree, NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, TieredSmt, ValuePath, Vec,
4};
5use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
6use core::borrow::Borrow;
7
8#[cfg(test)]
9mod tests;
10
11#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
12pub struct Node {
13 left: RpoDigest,
14 right: RpoDigest,
15}
16
17/// An in-memory data store for Merkelized data.
18///
19/// This is a in memory data store for Merkle trees, this store allows all the nodes of multiple
20/// trees to live as long as necessary and without duplication, this allows the implementation of
21/// space efficient persistent data structures.
22///
23/// Example usage:
24///
25/// ```rust
26/// # use miden_crypto::{ZERO, Felt, Word};
27/// # use miden_crypto::merkle::{NodeIndex, MerkleStore, MerkleTree};
28/// # use miden_crypto::hash::rpo::Rpo256;
29/// # const fn int_to_node(value: u64) -> Word {
30/// # [Felt::new(value), ZERO, ZERO, ZERO]
31/// # }
32/// # let A = int_to_node(1);
33/// # let B = int_to_node(2);
34/// # let C = int_to_node(3);
35/// # let D = int_to_node(4);
36/// # let E = int_to_node(5);
37/// # let F = int_to_node(6);
38/// # let G = int_to_node(7);
39/// # let H0 = int_to_node(8);
40/// # let H1 = int_to_node(9);
41/// # let T0 = MerkleTree::new([A, B, C, D, E, F, G, H0].to_vec()).expect("even number of leaves provided");
42/// # let T1 = MerkleTree::new([A, B, C, D, E, F, G, H1].to_vec()).expect("even number of leaves provided");
43/// # let ROOT0 = T0.root();
44/// # let ROOT1 = T1.root();
45/// let mut store = MerkleStore::new();
46///
47/// // the store is initialized with the SMT empty nodes
48/// assert_eq!(store.num_internal_nodes(), 255);
49///
50/// let tree1 = MerkleTree::new(vec![A, B, C, D, E, F, G, H0]).unwrap();
51/// let tree2 = MerkleTree::new(vec![A, B, C, D, E, F, G, H1]).unwrap();
52///
53/// // populates the store with two merkle trees, common nodes are shared
54/// store
55/// .extend(tree1.inner_nodes())
56/// .extend(tree2.inner_nodes());
57///
58/// // every leaf except the last are the same
59/// for i in 0..7 {
60/// let idx0 = NodeIndex::new(3, i).unwrap();
61/// let d0 = store.get_node(ROOT0, idx0).unwrap();
62/// let idx1 = NodeIndex::new(3, i).unwrap();
63/// let d1 = store.get_node(ROOT1, idx1).unwrap();
64/// assert_eq!(d0, d1, "Both trees have the same leaf at pos {i}");
65/// }
66///
67/// // The leafs A-B-C-D are the same for both trees, so are their 2 immediate parents
68/// for i in 0..4 {
69/// let idx0 = NodeIndex::new(3, i).unwrap();
70/// let d0 = store.get_path(ROOT0, idx0).unwrap();
71/// let idx1 = NodeIndex::new(3, i).unwrap();
72/// let d1 = store.get_path(ROOT1, idx1).unwrap();
73/// assert_eq!(d0.path[0..2], d1.path[0..2], "Both sub-trees are equal up to two levels");
74/// }
75///
76/// // Common internal nodes are shared, the two added trees have a total of 30, but the store has
77/// // only 10 new entries, corresponding to the 10 unique internal nodes of these trees.
78/// assert_eq!(store.num_internal_nodes() - 255, 10);
79/// ```
80#[derive(Debug, Clone, Eq, PartialEq)]
81pub struct MerkleStore {
82 nodes: BTreeMap<RpoDigest, Node>,
83}
84
85impl Default for MerkleStore {
86 fn default() -> Self {
87 Self::new()
88 }
89}
90
91impl MerkleStore {
92 // CONSTRUCTORS
93 // --------------------------------------------------------------------------------------------
94
95 /// Creates an empty `MerkleStore` instance.
96 pub fn new() -> MerkleStore {
97 // pre-populate the store with the empty hashes
98 let subtrees = EmptySubtreeRoots::empty_hashes(255);
99 let nodes = subtrees
100 .iter()
101 .rev()
102 .copied()
103 .zip(subtrees.iter().rev().skip(1).copied())
104 .map(|(child, parent)| {
105 (
106 parent,
107 Node {
108 left: child,
109 right: child,
110 },
111 )
112 })
113 .collect();
114
115 MerkleStore { nodes }
116 }
117
118 // PUBLIC ACCESSORS
119 // --------------------------------------------------------------------------------------------
120
121 /// Return a count of the non-leaf nodes in the store.
122 pub fn num_internal_nodes(&self) -> usize {
123 self.nodes.len()
124 }
125
126 /// Returns the node at `index` rooted on the tree `root`.
127 ///
128 /// # Errors
129 ///
130 /// This method can return the following errors:
131 /// - `RootNotInStore` if the `root` is not present in the store.
132 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
133 pub fn get_node(&self, root: RpoDigest, index: NodeIndex) -> Result<RpoDigest, MerkleError> {
134 let mut hash = root;
135
136 // corner case: check the root is in the store when called with index `NodeIndex::root()`
137 self.nodes.get(&hash).ok_or(MerkleError::RootNotInStore(hash))?;
138
139 for i in (0..index.depth()).rev() {
140 let node = self.nodes.get(&hash).ok_or(MerkleError::NodeNotInStore(hash, index))?;
141
142 let bit = (index.value() >> i) & 1;
143 hash = if bit == 0 { node.left } else { node.right }
144 }
145
146 Ok(hash)
147 }
148
149 /// Returns the node at the specified `index` and its opening to the `root`.
150 ///
151 /// The path starts at the sibling of the target leaf.
152 ///
153 /// # Errors
154 /// This method can return the following errors:
155 /// - `RootNotInStore` if the `root` is not present in the store.
156 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in 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 {
181 value: hash,
182 path: MerklePath::new(path),
183 })
184 }
185
186 /// Reconstructs a path from the root until a leaf or empty node and returns its depth.
187 ///
188 /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
189 /// from `root`. The maximum value the argument accepts is [u64::BITS].
190 ///
191 /// The traversed path from leaf to root will start at the least significant bit of `index`,
192 /// and will be executed for `tree_depth` bits.
193 ///
194 /// # Errors
195 /// Will return an error if:
196 /// - The provided root is not found.
197 /// - The path from the root continues to a depth greater than `tree_depth`.
198 /// - The provided `tree_depth` is greater than `64.
199 /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
200 /// information, check [NodeIndex::new].
201 pub fn get_leaf_depth(
202 &self,
203 root: RpoDigest,
204 tree_depth: u8,
205 index: u64,
206 ) -> Result<u8, MerkleError> {
207 // validate depth and index
208 if tree_depth > 64 {
209 return Err(MerkleError::DepthTooBig(tree_depth as u64));
210 }
211 NodeIndex::new(tree_depth, index)?;
212
213 // it's not illegal to have a maximum depth of `0`; we should just return the root in that
214 // case. this check will simplify the implementation as we could overflow bits for depth
215 // `0`.
216 if tree_depth == 0 {
217 return Ok(0);
218 }
219
220 // check if the root exists, providing the proper error report if it doesn't
221 let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
222 let mut hash = root;
223 if !self.nodes.contains_key(&hash) {
224 return Err(MerkleError::RootNotInStore(hash));
225 }
226
227 // we traverse from root to leaf, so the path is reversed
228 let mut path = (index << (64 - tree_depth)).reverse_bits();
229
230 // iterate every depth and reconstruct the path from root to leaf
231 for depth in 0..tree_depth {
232 // we short-circuit if an empty node has been found
233 if hash == empty[depth as usize] {
234 return Ok(depth);
235 }
236
237 // fetch the children pair, mapped by its parent hash
238 let children = match self.nodes.get(&hash) {
239 Some(node) => node,
240 None => return Ok(depth),
241 };
242
243 // traverse down
244 hash = if path & 1 == 0 { children.left } else { children.right };
245 path >>= 1;
246 }
247
248 // at max depth assert it doesn't have sub-trees
249 if self.nodes.contains_key(&hash) {
250 return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
251 }
252
253 // depleted bits; return max depth
254 Ok(tree_depth)
255 }
256
257 // DATA EXTRACTORS
258 // --------------------------------------------------------------------------------------------
259
260 /// Returns a subset of this Merkle store such that the returned Merkle store contains all
261 /// nodes which are descendants of the specified roots.
262 ///
263 /// The roots for which no descendants exist in this Merkle store are ignored.
264 pub fn subset<I, R>(&self, roots: I) -> MerkleStore
265 where
266 I: Iterator<Item = R>,
267 R: Borrow<RpoDigest>,
268 {
269 let mut store = MerkleStore::new();
270 for root in roots {
271 let root = *root.borrow();
272 store.clone_tree_from(root, self);
273 }
274 store
275 }
276
277 /// Iterator over the inner nodes of the [MerkleStore].
278 pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
279 self.nodes.iter().map(|(r, n)| InnerNodeInfo {
280 value: *r,
281 left: n.left,
282 right: n.right,
283 })
284 }
285
286 // STATE MUTATORS
287 // --------------------------------------------------------------------------------------------
288
289 /// Adds a sequence of nodes yielded by the provided iterator into the store.
290 pub fn extend<I>(&mut self, iter: I) -> &mut MerkleStore
291 where
292 I: Iterator<Item = InnerNodeInfo>,
293 {
294 for node in iter {
295 let value: RpoDigest = node.value;
296 let left: RpoDigest = node.left;
297 let right: RpoDigest = node.right;
298
299 debug_assert_eq!(Rpo256::merge(&[left, right]), value);
300 self.nodes.insert(value, Node { left, right });
301 }
302
303 self
304 }
305
306 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
307 /// new root.
308 ///
309 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
310 /// include all the nodes into the store.
311 pub fn add_merkle_path(
312 &mut self,
313 index: u64,
314 node: RpoDigest,
315 path: MerklePath,
316 ) -> Result<RpoDigest, MerkleError> {
317 let root = path.inner_nodes(index, node)?.fold(RpoDigest::default(), |_, node| {
318 let value: RpoDigest = node.value;
319 let left: RpoDigest = node.left;
320 let right: RpoDigest = node.right;
321
322 debug_assert_eq!(Rpo256::merge(&[left, right]), value);
323 self.nodes.insert(value, Node { left, right });
324
325 node.value
326 });
327 Ok(root)
328 }
329
330 /// Adds all the nodes of multiple Merkle paths into the store.
331 ///
332 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
333 /// into the store.
334 ///
335 /// For further reference, check [MerkleStore::add_merkle_path].
336 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<(), MerkleError>
337 where
338 I: IntoIterator<Item = (u64, RpoDigest, MerklePath)>,
339 {
340 for (index_value, node, path) in paths.into_iter() {
341 self.add_merkle_path(index_value, node, path)?;
342 }
343 Ok(())
344 }
345
346 /// Appends the provided [MerklePathSet] into the store.
347 ///
348 /// For further reference, check [MerkleStore::add_merkle_path].
349 pub fn add_merkle_path_set(
350 &mut self,
351 path_set: &MerklePathSet,
352 ) -> Result<RpoDigest, MerkleError> {
353 let root = path_set.root();
354 for (index, path) in path_set.to_paths() {
355 self.add_merkle_path(index, path.value, path.path)?;
356 }
357 Ok(root)
358 }
359
360 /// Sets a node to `value`.
361 ///
362 /// # Errors
363 ///
364 /// This method can return the following errors:
365 /// - `RootNotInStore` if the `root` is not present in the store.
366 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
367 pub fn set_node(
368 &mut self,
369 mut root: RpoDigest,
370 index: NodeIndex,
371 value: RpoDigest,
372 ) -> Result<RootPath, MerkleError> {
373 let node = value;
374 let ValuePath { value, path } = self.get_path(root, index)?;
375
376 // performs the update only if the node value differs from the opening
377 if node != value {
378 root = self.add_merkle_path(index.value(), node, path.clone())?;
379 }
380
381 Ok(RootPath { root, path })
382 }
383
384 /// Merges two elements and adds the resulting node into the store.
385 ///
386 /// Merges arbitrary values. They may be leafs, nodes, or a mixture of both.
387 pub fn merge_roots(
388 &mut self,
389 left_root: RpoDigest,
390 right_root: RpoDigest,
391 ) -> Result<RpoDigest, MerkleError> {
392 let parent = Rpo256::merge(&[left_root, right_root]);
393 self.nodes.insert(
394 parent,
395 Node {
396 left: left_root,
397 right: right_root,
398 },
399 );
400
401 Ok(parent)
402 }
403
404 // HELPER METHODS
405 // --------------------------------------------------------------------------------------------
406
407 /// Recursively clones a tree with the specified root from the specified source into self.
408 ///
409 /// If the source store does not contain a tree with the specified root, this is a noop.
410 fn clone_tree_from(&mut self, root: RpoDigest, source: &Self) {
411 // process the node only if it is in the source
412 if let Some(node) = source.nodes.get(&root) {
413 // if the node has already been inserted, no need to process it further as all of its
414 // descendants should be already cloned from the source store
415 if self.nodes.insert(root, *node).is_none() {
416 self.clone_tree_from(node.left, source);
417 self.clone_tree_from(node.right, source);
418 }
419 }
420 }
421}
422
423// CONVERSIONS
424// ================================================================================================
425
426impl From<&MerkleTree> for MerkleStore {
427 fn from(value: &MerkleTree) -> Self {
428 let mut store = MerkleStore::new();
429 store.extend(value.inner_nodes());
430 store
431 }
432}
433
434impl From<&SimpleSmt> for MerkleStore {
435 fn from(value: &SimpleSmt) -> Self {
436 let mut store = MerkleStore::new();
437 store.extend(value.inner_nodes());
438 store
439 }
440}
441
442impl From<&Mmr> for MerkleStore {
443 fn from(value: &Mmr) -> Self {
444 let mut store = MerkleStore::new();
445 store.extend(value.inner_nodes());
446 store
447 }
448}
449
450impl From<&TieredSmt> for MerkleStore {
451 fn from(value: &TieredSmt) -> Self {
452 let mut store = MerkleStore::new();
453 store.extend(value.inner_nodes());
454 store
455 }
456}
457
458impl FromIterator<InnerNodeInfo> for MerkleStore {
459 fn from_iter<T: IntoIterator<Item = InnerNodeInfo>>(iter: T) -> Self {
460 let mut store = MerkleStore::new();
461 store.extend(iter.into_iter());
462 store
463 }
464}
465
466// ITERATORS
467// ================================================================================================
468
469impl Extend<InnerNodeInfo> for MerkleStore {
470 fn extend<T: IntoIterator<Item = InnerNodeInfo>>(&mut self, iter: T) {
471 self.extend(iter.into_iter());
472 }
473}
474
475// SERIALIZATION
476// ================================================================================================
477
478impl Serializable for Node {
479 fn write_into<W: ByteWriter>(&self, target: &mut W) {
480 self.left.write_into(target);
481 self.right.write_into(target);
482 }
483}
484
485impl Deserializable for Node {
486 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
487 let left = RpoDigest::read_from(source)?;
488 let right = RpoDigest::read_from(source)?;
489 Ok(Node { left, right })
490 }
491}
492
493impl Serializable for MerkleStore {
494 fn write_into<W: ByteWriter>(&self, target: &mut W) {
495 target.write_u64(self.nodes.len() as u64);
496
497 for (k, v) in self.nodes.iter() {
498 k.write_into(target);
499 v.write_into(target);
500 }
501 }
502}
503
504impl Deserializable for MerkleStore {
505 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
506 let len = source.read_u64()?;
507 let mut nodes: BTreeMap<RpoDigest, Node> = BTreeMap::new();
508
509 for _ in 0..len {
510 let key = RpoDigest::read_from(source)?;
511 let value = Node::read_from(source)?;
512 nodes.insert(key, value);
513 }
514
515 Ok(MerkleStore { nodes })
516 }
517}