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