miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.rsRaw
1use super::mmr::{Mmr, MmrPeaks};
2use super::{
3 BTreeMap, EmptySubtreeRoots, MerkleError, MerklePath, MerklePathSet, MerkleTree, NodeIndex,
4 RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
5};
6use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
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 Merkle-lized 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/// // populates the store with two merkle trees, common nodes are shared
51/// store.add_merkle_tree([A, B, C, D, E, F, G, H0]);
52/// store.add_merkle_tree([A, B, C, D, E, F, G, H1]);
53///
54/// // every leaf except the last are the same
55/// for i in 0..7 {
56/// let idx0 = NodeIndex::new(3, i).unwrap();
57/// let d0 = store.get_node(ROOT0, idx0).unwrap();
58/// let idx1 = NodeIndex::new(3, i).unwrap();
59/// let d1 = store.get_node(ROOT1, idx1).unwrap();
60/// assert_eq!(d0, d1, "Both trees have the same leaf at pos {i}");
61/// }
62///
63/// // The leafs A-B-C-D are the same for both trees, so are their 2 immediate parents
64/// for i in 0..4 {
65/// let idx0 = NodeIndex::new(3, i).unwrap();
66/// let d0 = store.get_path(ROOT0, idx0).unwrap();
67/// let idx1 = NodeIndex::new(3, i).unwrap();
68/// let d1 = store.get_path(ROOT1, idx1).unwrap();
69/// assert_eq!(d0.path[0..2], d1.path[0..2], "Both sub-trees are equal up to two levels");
70/// }
71///
72/// // Common internal nodes are shared, the two added trees have a total of 30, but the store has
73/// // only 10 new entries, corresponding to the 10 unique internal nodes of these trees.
74/// assert_eq!(store.num_internal_nodes() - 255, 10);
75/// ```
76#[derive(Debug, Clone, Eq, PartialEq)]
77pub struct MerkleStore {
78 nodes: BTreeMap<RpoDigest, Node>,
79}
80
81impl Default for MerkleStore {
82 fn default() -> Self {
83 Self::new()
84 }
85}
86
87impl MerkleStore {
88 // CONSTRUCTORS
89 // --------------------------------------------------------------------------------------------
90
91 /// Creates an empty `MerkleStore` instance.
92 pub fn new() -> MerkleStore {
93 // pre-populate the store with the empty hashes
94 let subtrees = EmptySubtreeRoots::empty_hashes(255);
95 let nodes = subtrees
96 .iter()
97 .rev()
98 .copied()
99 .zip(subtrees.iter().rev().skip(1).copied())
100 .map(|(child, parent)| {
101 (
102 parent,
103 Node {
104 left: child,
105 right: child,
106 },
107 )
108 })
109 .collect();
110
111 MerkleStore { nodes }
112 }
113
114 /// Appends the provided merkle tree represented by its `leaves` to the set.
115 pub fn with_merkle_tree<I>(mut self, leaves: I) -> Result<Self, MerkleError>
116 where
117 I: IntoIterator<Item = Word>,
118 {
119 self.add_merkle_tree(leaves)?;
120 Ok(self)
121 }
122
123 /// Appends the provided Sparse Merkle tree represented by its `entries` to the set.
124 ///
125 /// For more information, check [MerkleStore::add_sparse_merkle_tree].
126 pub fn with_sparse_merkle_tree<R, I>(
127 mut self,
128 depth: u8,
129 entries: R,
130 ) -> Result<Self, MerkleError>
131 where
132 R: IntoIterator<IntoIter = I>,
133 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
134 {
135 self.add_sparse_merkle_tree(depth, entries)?;
136 Ok(self)
137 }
138
139 /// Appends the provided merkle path set.
140 pub fn with_merkle_path(
141 mut self,
142 index_value: u64,
143 node: Word,
144 path: MerklePath,
145 ) -> Result<Self, MerkleError> {
146 self.add_merkle_path(index_value, node, path)?;
147 Ok(self)
148 }
149
150 /// Appends the provided merkle path set.
151 pub fn with_merkle_paths<I>(mut self, paths: I) -> Result<Self, MerkleError>
152 where
153 I: IntoIterator<Item = (u64, Word, MerklePath)>,
154 {
155 self.add_merkle_paths(paths)?;
156 Ok(self)
157 }
158
159 /// Appends the provided [Mmr] represented by its `leaves` to the set.
160 pub fn with_mmr<I>(mut self, leaves: I) -> Result<Self, MerkleError>
161 where
162 I: IntoIterator<Item = Word>,
163 {
164 self.add_mmr(leaves)?;
165 Ok(self)
166 }
167
168 // PUBLIC ACCESSORS
169 // --------------------------------------------------------------------------------------------
170
171 /// Return a count of the non-leaf nodes in the store.
172 pub fn num_internal_nodes(&self) -> usize {
173 self.nodes.len()
174 }
175
176 /// Returns the node at `index` rooted on the tree `root`.
177 ///
178 /// # Errors
179 ///
180 /// This method can return the following errors:
181 /// - `RootNotInStore` if the `root` is not present in the store.
182 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
183 pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
184 let mut hash: RpoDigest = root.into();
185
186 // corner case: check the root is in the store when called with index `NodeIndex::root()`
187 self.nodes
188 .get(&hash)
189 .ok_or(MerkleError::RootNotInStore(hash.into()))?;
190
191 for i in (0..index.depth()).rev() {
192 let node = self
193 .nodes
194 .get(&hash)
195 .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
196
197 let bit = (index.value() >> i) & 1;
198 hash = if bit == 0 { node.left } else { node.right }
199 }
200
201 Ok(hash.into())
202 }
203
204 /// Returns the node at the specified `index` and its opening to the `root`.
205 ///
206 /// The path starts at the sibling of the target leaf.
207 ///
208 /// # Errors
209 ///
210 /// This method can return the following errors:
211 /// - `RootNotInStore` if the `root` is not present in the store.
212 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
213 pub fn get_path(&self, root: Word, index: NodeIndex) -> Result<ValuePath, MerkleError> {
214 let mut hash: RpoDigest = root.into();
215 let mut path = Vec::with_capacity(index.depth().into());
216
217 // corner case: check the root is in the store when called with index `NodeIndex::root()`
218 self.nodes
219 .get(&hash)
220 .ok_or(MerkleError::RootNotInStore(hash.into()))?;
221
222 for i in (0..index.depth()).rev() {
223 let node = self
224 .nodes
225 .get(&hash)
226 .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
227
228 let bit = (index.value() >> i) & 1;
229 hash = if bit == 0 {
230 path.push(node.right.into());
231 node.left
232 } else {
233 path.push(node.left.into());
234 node.right
235 }
236 }
237
238 // the path is computed from root to leaf, so it must be reversed
239 path.reverse();
240
241 Ok(ValuePath {
242 value: hash.into(),
243 path: MerklePath::new(path),
244 })
245 }
246
247 /// Reconstructs a path from the root until a leaf or empty node and returns its depth.
248 ///
249 /// The `tree_depth` parameter defines up to which depth the tree will be traversed, starting
250 /// from `root`. The maximum value the argument accepts is [u64::BITS].
251 ///
252 /// The traversed path from leaf to root will start at the least significant bit of `index`,
253 /// and will be executed for `tree_depth` bits.
254 ///
255 /// # Errors
256 /// Will return an error if:
257 /// - The provided root is not found.
258 /// - The path from the root continues to a depth greater than `tree_depth`.
259 /// - The provided `tree_depth` is greater than `64.
260 /// - The provided `index` is not valid for a depth equivalent to `tree_depth`. For more
261 /// information, check [NodeIndex::new].
262 pub fn get_leaf_depth(
263 &self,
264 root: Word,
265 tree_depth: u8,
266 index: u64,
267 ) -> Result<u8, MerkleError> {
268 // validate depth and index
269 if tree_depth > 64 {
270 return Err(MerkleError::DepthTooBig(tree_depth as u64));
271 }
272 NodeIndex::new(tree_depth, index)?;
273
274 // it's not illegal to have a maximum depth of `0`; we should just return the root in that
275 // case. this check will simplify the implementation as we could overflow bits for depth
276 // `0`.
277 if tree_depth == 0 {
278 return Ok(0);
279 }
280
281 // check if the root exists, providing the proper error report if it doesn't
282 let empty = EmptySubtreeRoots::empty_hashes(tree_depth);
283 let mut hash: RpoDigest = root.into();
284 if !self.nodes.contains_key(&hash) {
285 return Err(MerkleError::RootNotInStore(hash.into()));
286 }
287
288 // we traverse from root to leaf, so the path is reversed
289 let mut path = (index << (64 - tree_depth)).reverse_bits();
290
291 // iterate every depth and reconstruct the path from root to leaf
292 for depth in 0..tree_depth {
293 // we short-circuit if an empty node has been found
294 if hash == empty[depth as usize] {
295 return Ok(depth);
296 }
297
298 // fetch the children pair, mapped by its parent hash
299 let children = match self.nodes.get(&hash) {
300 Some(node) => node,
301 None => return Ok(depth),
302 };
303
304 // traverse down
305 hash = if path & 1 == 0 {
306 children.left
307 } else {
308 children.right
309 };
310 path >>= 1;
311 }
312
313 // at max depth assert it doesn't have sub-trees
314 if self.nodes.contains_key(&hash) {
315 return Err(MerkleError::DepthTooBig(tree_depth as u64 + 1));
316 }
317
318 // depleted bits; return max depth
319 Ok(tree_depth)
320 }
321
322 // STATE MUTATORS
323 // --------------------------------------------------------------------------------------------
324
325 /// Adds all the nodes of a Merkle tree represented by `leaves`.
326 ///
327 /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
328 /// store.
329 ///
330 /// # Errors
331 ///
332 /// This method may return the following errors:
333 /// - `DepthTooSmall` if leaves is empty or contains only 1 element
334 /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
335 pub fn add_merkle_tree<I>(&mut self, leaves: I) -> Result<Word, MerkleError>
336 where
337 I: IntoIterator<Item = Word>,
338 {
339 let leaves: Vec<_> = leaves.into_iter().collect();
340 if leaves.len() < 2 {
341 return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
342 }
343
344 let tree = MerkleTree::new(leaves)?;
345 for node in tree.inner_nodes() {
346 self.nodes.insert(
347 node.value.into(),
348 Node {
349 left: node.left.into(),
350 right: node.right.into(),
351 },
352 );
353 }
354
355 Ok(tree.root())
356 }
357
358 /// Adds a Sparse Merkle tree defined by the specified `entries` to the store, and returns the
359 /// root of the added tree.
360 ///
361 /// The entries are expected to contain tuples of `(index, node)` describing nodes in the tree
362 /// at `depth`.
363 ///
364 /// # Errors
365 /// Returns an error if the provided `depth` is greater than [SimpleSmt::MAX_DEPTH].
366 pub fn add_sparse_merkle_tree<R, I>(
367 &mut self,
368 depth: u8,
369 entries: R,
370 ) -> Result<Word, MerkleError>
371 where
372 R: IntoIterator<IntoIter = I>,
373 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
374 {
375 let smt = SimpleSmt::new(depth)?.with_leaves(entries)?;
376 for node in smt.inner_nodes() {
377 self.nodes.insert(
378 node.value.into(),
379 Node {
380 left: node.left.into(),
381 right: node.right.into(),
382 },
383 );
384 }
385
386 Ok(smt.root())
387 }
388
389 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
390 /// new root.
391 ///
392 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
393 /// include all the nodes into the store.
394 pub fn add_merkle_path(
395 &mut self,
396 index_value: u64,
397 mut node: Word,
398 path: MerklePath,
399 ) -> Result<Word, MerkleError> {
400 let mut index = NodeIndex::new(path.len() as u8, index_value)?;
401
402 for sibling in path {
403 let (left, right) = match index.is_value_odd() {
404 true => (sibling, node),
405 false => (node, sibling),
406 };
407 let parent = Rpo256::merge(&[left.into(), right.into()]);
408 self.nodes.insert(
409 parent,
410 Node {
411 left: left.into(),
412 right: right.into(),
413 },
414 );
415
416 index.move_up();
417 node = parent.into();
418 }
419
420 Ok(node)
421 }
422
423 /// Adds all the nodes of multiple Merkle paths into the store.
424 ///
425 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
426 /// into the store.
427 ///
428 /// For further reference, check [MerkleStore::add_merkle_path].
429 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<(), MerkleError>
430 where
431 I: IntoIterator<Item = (u64, Word, MerklePath)>,
432 {
433 for (index_value, node, path) in paths.into_iter() {
434 self.add_merkle_path(index_value, node, path)?;
435 }
436 Ok(())
437 }
438
439 /// Appends the provided [MerklePathSet] into the store.
440 ///
441 /// For further reference, check [MerkleStore::add_merkle_path].
442 pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result<Word, MerkleError> {
443 let root = path_set.root();
444 for (index, path) in path_set.to_paths() {
445 self.add_merkle_path(index, path.value, path.path)?;
446 }
447 Ok(root)
448 }
449
450 /// Appends the provided [Mmr] into the store.
451 pub fn add_mmr<I>(&mut self, leaves: I) -> Result<MmrPeaks, MerkleError>
452 where
453 I: IntoIterator<Item = Word>,
454 {
455 let mmr = Mmr::from(leaves);
456 for node in mmr.inner_nodes() {
457 self.nodes.insert(
458 node.value.into(),
459 Node {
460 left: node.left.into(),
461 right: node.right.into(),
462 },
463 );
464 }
465
466 Ok(mmr.accumulator())
467 }
468
469 /// Sets a node to `value`.
470 ///
471 /// # Errors
472 ///
473 /// This method can return the following errors:
474 /// - `RootNotInStore` if the `root` is not present in the store.
475 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
476 pub fn set_node(
477 &mut self,
478 mut root: Word,
479 index: NodeIndex,
480 value: Word,
481 ) -> Result<RootPath, MerkleError> {
482 let node = value;
483 let ValuePath { value, path } = self.get_path(root, index)?;
484
485 // performs the update only if the node value differs from the opening
486 if node != value {
487 root = self.add_merkle_path(index.value(), node, path.clone())?;
488 }
489
490 Ok(RootPath { root, path })
491 }
492
493 pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
494 let root1: RpoDigest = root1.into();
495 let root2: RpoDigest = root2.into();
496
497 if !self.nodes.contains_key(&root1) {
498 Err(MerkleError::NodeNotInStore(root1.into(), NodeIndex::root()))
499 } else if !self.nodes.contains_key(&root1) {
500 Err(MerkleError::NodeNotInStore(root2.into(), NodeIndex::root()))
501 } else {
502 let parent: Word = Rpo256::merge(&[root1, root2]).into();
503 self.nodes.insert(
504 parent.into(),
505 Node {
506 left: root1,
507 right: root2,
508 },
509 );
510
511 Ok(parent)
512 }
513 }
514}
515
516// SERIALIZATION
517// ================================================================================================
518
519impl Serializable for Node {
520 fn write_into<W: ByteWriter>(&self, target: &mut W) {
521 self.left.write_into(target);
522 self.right.write_into(target);
523 }
524}
525
526impl Deserializable for Node {
527 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
528 let left = RpoDigest::read_from(source)?;
529 let right = RpoDigest::read_from(source)?;
530 Ok(Node { left, right })
531 }
532}
533
534impl Serializable for MerkleStore {
535 fn write_into<W: ByteWriter>(&self, target: &mut W) {
536 target.write_u64(self.nodes.len() as u64);
537
538 for (k, v) in self.nodes.iter() {
539 k.write_into(target);
540 v.write_into(target);
541 }
542 }
543}
544
545impl Deserializable for MerkleStore {
546 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
547 let len = source.read_u64()?;
548 let mut nodes: BTreeMap<RpoDigest, Node> = BTreeMap::new();
549
550 for _ in 0..len {
551 let key = RpoDigest::read_from(source)?;
552 let value = Node::read_from(source)?;
553 nodes.insert(key, value);
554 }
555
556 Ok(MerkleStore { nodes })
557 }
558}