miden-crypto/src/merkle/store.rs
src/merkle/store.rsRaw
1//! An in-memory data store for Merkle-lized data
2//!
3//! This is a in memory data store for Merkle trees, this store allows all the nodes of a tree
4//! (leaves or internal) to live as long as necessary and without duplication, this allows the
5//! implementation of efficient persistent data structures
6use super::{
7 BTreeMap, BTreeSet, EmptySubtreeRoots, MerkleError, MerklePath, MerkleTree, NodeIndex, Rpo256,
8 RpoDigest, SimpleSmt, Vec, Word,
9};
10
11#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
12pub struct Node {
13 left: RpoDigest,
14 right: RpoDigest,
15}
16
17#[derive(Debug, Clone, Eq, PartialEq)]
18pub struct MerkleStore {
19 nodes: BTreeMap<RpoDigest, Node>,
20}
21
22impl Default for MerkleStore {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27
28impl MerkleStore {
29 // CONSTRUCTORS
30 // --------------------------------------------------------------------------------------------
31
32 /// Creates an empty `MerkleStore` instance.
33 pub fn new() -> MerkleStore {
34 // pre-populate the store with the empty hashes
35 let subtrees = EmptySubtreeRoots::empty_hashes(64);
36 let nodes = subtrees
37 .iter()
38 .copied()
39 .zip(subtrees.iter().skip(1).copied())
40 .map(|(child, parent)| {
41 (
42 parent,
43 Node {
44 left: child,
45 right: child,
46 },
47 )
48 })
49 .collect();
50
51 MerkleStore { nodes }
52 }
53
54 /// Appends the provided merkle tree represented by its `leaves` to the set.
55 pub fn with_merkle_tree<I>(mut self, leaves: I) -> Result<Self, MerkleError>
56 where
57 I: IntoIterator<Item = Word>,
58 {
59 self.add_merkle_tree(leaves)?;
60 Ok(self)
61 }
62
63 /// Appends the provided sparse merkle tree represented by its `entries` to the set.
64 pub fn with_sparse_merkle_tree<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
65 where
66 R: IntoIterator<IntoIter = I>,
67 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
68 {
69 self.add_sparse_merkle_tree(entries)?;
70 Ok(self)
71 }
72
73 /// Appends the provided merkle path set.
74 pub fn with_merkle_path(
75 mut self,
76 index_value: u64,
77 node: Word,
78 path: MerklePath,
79 ) -> Result<Self, MerkleError> {
80 self.add_merkle_path(index_value, node, path)?;
81 Ok(self)
82 }
83
84 // STATE MUTATORS
85 // --------------------------------------------------------------------------------------------
86
87 /// Adds all the nodes of a Merkle tree represented by `leaves`.
88 ///
89 /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
90 /// storage.
91 ///
92 /// # Errors
93 ///
94 /// This method may return the following errors:
95 /// - `DepthTooSmall` if leaves is empty or contains only 1 element
96 /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
97 pub fn add_merkle_tree<I>(&mut self, leaves: I) -> Result<Word, MerkleError>
98 where
99 I: IntoIterator<Item = Word>,
100 {
101 let leaves: Vec<_> = leaves.into_iter().collect();
102 let layers = leaves.len().ilog2();
103 let tree = MerkleTree::new(leaves)?;
104
105 let mut depth = 0;
106 let mut parent_offset = 1;
107 let mut child_offset = 2;
108 while depth < layers {
109 let layer_size = 1usize << depth;
110 for _ in 0..layer_size {
111 // merkle tree is using level form representation, so left and right siblings are
112 // next to each other
113 let left = tree.nodes[child_offset];
114 let right = tree.nodes[child_offset + 1];
115 self.nodes.insert(
116 tree.nodes[parent_offset].into(),
117 Node {
118 left: left.into(),
119 right: right.into(),
120 },
121 );
122 parent_offset += 1;
123 child_offset += 2;
124 }
125 depth += 1;
126 }
127
128 Ok(tree.nodes[1])
129 }
130
131 /// Adds all the nodes of a Sparse Merkle tree represented by `entries`.
132 ///
133 /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into
134 /// the storage.
135 ///
136 /// # Errors
137 ///
138 /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`.
139 pub fn add_sparse_merkle_tree<R, I>(&mut self, entries: R) -> Result<Word, MerkleError>
140 where
141 R: IntoIterator<IntoIter = I>,
142 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
143 {
144 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?;
145 for branch in smt.store.branches.values() {
146 let parent = Rpo256::merge(&[branch.left, branch.right]);
147 self.nodes.insert(
148 parent,
149 Node {
150 left: branch.left,
151 right: branch.right,
152 },
153 );
154 }
155
156 Ok(smt.root())
157 }
158
159 /// Adds all the nodes of a Merkle path represented by `path`.
160 ///
161 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
162 /// include all the nodes into the storage.
163 pub fn add_merkle_path(
164 &mut self,
165 index_value: u64,
166 node: Word,
167 path: MerklePath,
168 ) -> Result<Word, MerkleError> {
169 let mut node = node;
170 let mut index = NodeIndex::new(self.nodes.len() as u8, index_value);
171
172 for sibling in path {
173 let (left, right) = match index.is_value_odd() {
174 true => (sibling, node),
175 false => (node, sibling),
176 };
177 let parent = Rpo256::merge(&[left.into(), right.into()]);
178 self.nodes.insert(
179 parent,
180 Node {
181 left: left.into(),
182 right: right.into(),
183 },
184 );
185
186 index.move_up();
187 node = parent.into();
188 }
189
190 Ok(node)
191 }
192
193 /// Adds all the nodes of multiple Merkle paths into the store.
194 ///
195 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
196 /// into the storage.
197 ///
198 /// # Errors
199 ///
200 /// Every path must resolve to the same root, otherwise this will return an `ConflictingRoots`
201 /// error.
202 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<Word, MerkleError>
203 where
204 I: IntoIterator<Item = (u64, Word, MerklePath)>,
205 {
206 let paths: Vec<(u64, Word, MerklePath)> = paths.into_iter().collect();
207
208 let roots: BTreeSet<RpoDigest> = paths
209 .iter()
210 .map(|(index, node, path)| path.compute_root(*index, *node).into())
211 .collect();
212
213 if roots.len() != 1 {
214 return Err(MerkleError::ConflictingRoots(
215 roots.iter().map(|v| Word::from(*v)).collect(),
216 ));
217 }
218
219 for (index_value, node, path) in paths {
220 self.add_merkle_path(index_value, node, path)?;
221 }
222
223 // Returns the parent of the last paths (assumes all paths have the same parent) or empty
224 // The length of unique_roots is checked above, so this wont panic
225 Ok(roots.iter().next().unwrap().into())
226 }
227
228 // PUBLIC ACCESSORS
229 // --------------------------------------------------------------------------------------------
230
231 /// Returns the node at `index` rooted on the tree `root`.
232 ///
233 /// # Errors
234 ///
235 /// This will return `NodeNotInStorage` if the element is not present in the store.
236 pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
237 let mut hash: RpoDigest = root.into();
238
239 // Check the root is in the storage when called with `NodeIndex::root()`
240 self.nodes
241 .get(&hash)
242 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
243
244 for bit in index.bit_iterator().rev() {
245 let node = self
246 .nodes
247 .get(&hash)
248 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
249 hash = if bit { node.right } else { node.left }
250 }
251
252 Ok(hash.into())
253 }
254
255 /// Returns the path for the node at `index` rooted on the tree `root`.
256 ///
257 /// # Errors
258 ///
259 /// This will return `NodeNotInStorage` if the element is not present in the store.
260 pub fn get_path(
261 &self,
262 root: Word,
263 index: NodeIndex,
264 ) -> Result<(Word, MerklePath), MerkleError> {
265 let mut hash: RpoDigest = root.into();
266 let mut path = Vec::new();
267 let node = RpoDigest::default();
268 for bit in index.bit_iterator() {
269 let node = self
270 .nodes
271 .get(&hash)
272 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
273
274 hash = if bit {
275 path.push(node.left.into());
276 node.right
277 } else {
278 path.push(node.right.into());
279 node.left
280 }
281 }
282
283 Ok((node.into(), MerklePath::new(path)))
284 }
285
286 // DATA MUTATORS
287 // --------------------------------------------------------------------------------------------
288
289 pub fn set_node(
290 &mut self,
291 root: Word,
292 index: NodeIndex,
293 value: Word,
294 ) -> Result<Word, MerkleError> {
295 let (current_node, path) = self.get_path(root, index)?;
296 if current_node != value {
297 self.add_merkle_path(index.value(), value, path)
298 } else {
299 Ok(root)
300 }
301 }
302
303 pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
304 let root1: RpoDigest = root1.into();
305 let root2: RpoDigest = root2.into();
306
307 if !self.nodes.contains_key(&root1) {
308 Err(MerkleError::NodeNotInStorage(
309 root1.into(),
310 NodeIndex::new(0, 0),
311 ))
312 } else if !self.nodes.contains_key(&root1) {
313 Err(MerkleError::NodeNotInStorage(
314 root2.into(),
315 NodeIndex::new(0, 0),
316 ))
317 } else {
318 let parent: Word = Rpo256::merge(&[root1, root2]).into();
319 self.nodes.insert(
320 parent.into(),
321 Node {
322 left: root1,
323 right: root2,
324 },
325 );
326
327 Ok(parent)
328 }
329 }
330}
331
332#[cfg(test)]
333mod test {
334 use super::{MerkleError, MerkleStore, MerkleTree, NodeIndex, SimpleSmt, Word};
335 use crate::merkle::int_to_node;
336 use crate::merkle::MerklePathSet;
337
338 const KEYS4: [u64; 4] = [0, 1, 2, 3];
339 const LEAVES4: [Word; 4] = [
340 int_to_node(1),
341 int_to_node(2),
342 int_to_node(3),
343 int_to_node(4),
344 ];
345
346 #[test]
347 fn test_add_merkle_tree() -> Result<(), MerkleError> {
348 let mut store = MerkleStore::default();
349
350 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
351 store.add_merkle_tree(LEAVES4.to_vec())?;
352
353 assert!(
354 store
355 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
356 .is_ok(),
357 "node 0 must be in the tree"
358 );
359 assert!(
360 store
361 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
362 .is_ok(),
363 "node 1 must be in the tree"
364 );
365 assert!(
366 store
367 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
368 .is_ok(),
369 "node 2 must be in the tree"
370 );
371 assert!(
372 store
373 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
374 .is_ok(),
375 "node 3 must be in the tree"
376 );
377
378 store
379 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
380 .expect("node 0 must be in tree");
381 store
382 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
383 .expect("node 1 must be in tree");
384 store
385 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
386 .expect("node 2 must be in tree");
387 store
388 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
389 .expect("node 3 must be in tree");
390
391 Ok(())
392 }
393
394 #[test]
395 fn test_get_invalid_node() {
396 let mut store = MerkleStore::default();
397 let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
398 store
399 .add_merkle_tree(LEAVES4.to_vec())
400 .expect("adding a merkle tree to the store must work");
401 let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3));
402 }
403
404 #[test]
405 fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
406 let mut store = MerkleStore::default();
407 let keys2: [u64; 2] = [0, 1];
408 let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
409 store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
410 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
411 .unwrap()
412 .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
413 .unwrap();
414
415 let idx = NodeIndex::new(1, 0);
416 assert_eq!(
417 store.get_node(smt.root(), idx).unwrap(),
418 smt.get_node(&idx).unwrap()
419 );
420
421 Ok(())
422 }
423
424 #[test]
425 fn test_add_sparse_merkle_tree() -> Result<(), MerkleError> {
426 let mut store = MerkleStore::default();
427 store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
428
429 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
430 .unwrap()
431 .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
432 .unwrap();
433
434 let idx = NodeIndex::new(1, 0);
435 assert_eq!(
436 store.get_node(smt.root(), idx).unwrap(),
437 smt.get_node(&idx).unwrap()
438 );
439 let idx = NodeIndex::new(1, 1);
440 assert_eq!(
441 store.get_node(smt.root(), idx).unwrap(),
442 smt.get_node(&idx).unwrap()
443 );
444
445 Ok(())
446 }
447
448 #[test]
449 fn test_add_merkle_paths() -> Result<(), MerkleError> {
450 let mut store = MerkleStore::default();
451 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
452
453 let i0 = 0;
454 let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap();
455
456 let i1 = 1;
457 let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap();
458
459 let i2 = 2;
460 let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap();
461
462 let i3 = 3;
463 let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap();
464
465 let paths = [
466 (i0, LEAVES4[i0 as usize], p0),
467 (i1, LEAVES4[i1 as usize], p1),
468 (i2, LEAVES4[i2 as usize], p2),
469 (i3, LEAVES4[i3 as usize], p3),
470 ];
471
472 store
473 .add_merkle_paths(paths.clone())
474 .expect("the valid paths must work");
475
476 let set = MerklePathSet::new(3).with_paths(paths).unwrap();
477
478 assert_eq!(
479 set.get_node(NodeIndex::new(3, 0)).unwrap(),
480 store.get_node(set.root(), NodeIndex::new(2, 0b00)).unwrap(),
481 );
482 assert_eq!(
483 set.get_node(NodeIndex::new(3, 1)).unwrap(),
484 store.get_node(set.root(), NodeIndex::new(2, 0b01)).unwrap(),
485 );
486 assert_eq!(
487 set.get_node(NodeIndex::new(3, 2)).unwrap(),
488 store.get_node(set.root(), NodeIndex::new(2, 0b10)).unwrap(),
489 );
490 assert_eq!(
491 set.get_node(NodeIndex::new(3, 3)).unwrap(),
492 store.get_node(set.root(), NodeIndex::new(2, 0b11)).unwrap(),
493 );
494
495 Ok(())
496 }
497}