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 for bit in index.bit_iterator().rev() {
239 let node = self
240 .nodes
241 .get(&hash)
242 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
243 hash = if bit { node.right } else { node.left }
244 }
245
246 Ok(hash.into())
247 }
248
249 /// Returns the path for the node at `index` rooted on the tree `root`.
250 ///
251 /// # Errors
252 ///
253 /// This will return `NodeNotInStorage` if the element is not present in the store.
254 pub fn get_path(
255 &self,
256 root: Word,
257 index: NodeIndex,
258 ) -> Result<(Word, MerklePath), MerkleError> {
259 let mut hash: RpoDigest = root.into();
260 let mut path = Vec::new();
261 let node = RpoDigest::default();
262 for bit in index.bit_iterator() {
263 let node = self
264 .nodes
265 .get(&hash)
266 .ok_or(MerkleError::NodeNotInStorage(hash.into(), index))?;
267
268 hash = if bit {
269 path.push(node.left.into());
270 node.right
271 } else {
272 path.push(node.right.into());
273 node.left
274 }
275 }
276
277 Ok((node.into(), MerklePath::new(path)))
278 }
279
280 // DATA MUTATORS
281 // --------------------------------------------------------------------------------------------
282
283 pub fn set_node(
284 &mut self,
285 root: Word,
286 index: NodeIndex,
287 value: Word,
288 ) -> Result<Word, MerkleError> {
289 let (current_node, path) = self.get_path(root, index)?;
290 if current_node != value {
291 self.add_merkle_path(index.value(), value, path)
292 } else {
293 Ok(root)
294 }
295 }
296
297 pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
298 let root1: RpoDigest = root1.into();
299 let root2: RpoDigest = root2.into();
300
301 if !self.nodes.contains_key(&root1) {
302 Err(MerkleError::NodeNotInStorage(
303 root1.into(),
304 NodeIndex::new(0, 0),
305 ))
306 } else if !self.nodes.contains_key(&root1) {
307 Err(MerkleError::NodeNotInStorage(
308 root2.into(),
309 NodeIndex::new(0, 0),
310 ))
311 } else {
312 let parent: Word = Rpo256::merge(&[root1, root2]).into();
313 self.nodes.insert(
314 parent.into(),
315 Node {
316 left: root1,
317 right: root2,
318 },
319 );
320
321 Ok(parent)
322 }
323 }
324}
325
326#[cfg(test)]
327mod test {
328 use super::{MerkleError, MerkleStore, MerkleTree, NodeIndex, SimpleSmt, Word};
329 use crate::merkle::int_to_node;
330 use crate::merkle::MerklePathSet;
331
332 const KEYS4: [u64; 4] = [0, 1, 2, 3];
333 const LEAVES4: [Word; 4] = [
334 int_to_node(1),
335 int_to_node(2),
336 int_to_node(3),
337 int_to_node(4),
338 ];
339
340 #[test]
341 fn test_add_merkle_tree() -> Result<(), MerkleError> {
342 let mut store = MerkleStore::default();
343
344 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
345 store.add_merkle_tree(LEAVES4.to_vec())?;
346
347 assert!(
348 store
349 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
350 .is_ok(),
351 "node 0 must be in the tree"
352 );
353 assert!(
354 store
355 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
356 .is_ok(),
357 "node 1 must be in the tree"
358 );
359 assert!(
360 store
361 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
362 .is_ok(),
363 "node 2 must be in the tree"
364 );
365 assert!(
366 store
367 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
368 .is_ok(),
369 "node 3 must be in the tree"
370 );
371
372 store
373 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 0))
374 .expect("node 0 must be in tree");
375 store
376 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 1))
377 .expect("node 1 must be in tree");
378 store
379 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 2))
380 .expect("node 2 must be in tree");
381 store
382 .get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3))
383 .expect("node 3 must be in tree");
384
385 Ok(())
386 }
387
388 #[test]
389 fn test_get_node_returns_self_for_root() {
390 let store = MerkleStore::default();
391 let root_idx = NodeIndex::new(0, 0);
392
393 // the root does not need any lookups in the storage itself, so the value is just returned
394 assert_eq!(store.get_node(LEAVES4[0], root_idx).unwrap(), LEAVES4[0]);
395 assert_eq!(store.get_node(LEAVES4[1], root_idx).unwrap(), LEAVES4[1]);
396 assert_eq!(store.get_node(LEAVES4[2], root_idx).unwrap(), LEAVES4[2]);
397 assert_eq!(store.get_node(LEAVES4[3], root_idx).unwrap(), LEAVES4[3]);
398 }
399
400 #[test]
401 fn test_get_invalid_node() {
402 let mut store = MerkleStore::default();
403 let mtree = MerkleTree::new(LEAVES4.to_vec()).expect("creating a merkle tree must work");
404 store
405 .add_merkle_tree(LEAVES4.to_vec())
406 .expect("adding a merkle tree to the store must work");
407 let _ = store.get_node(mtree.root(), NodeIndex::new(mtree.depth(), 3));
408 }
409
410 #[test]
411 fn test_add_sparse_merkle_tree_one_level() -> Result<(), MerkleError> {
412 let mut store = MerkleStore::default();
413 let keys2: [u64; 2] = [0, 1];
414 let leaves2: [Word; 2] = [int_to_node(1), int_to_node(2)];
415 store.add_sparse_merkle_tree(keys2.into_iter().zip(leaves2.into_iter()))?;
416 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
417 .unwrap()
418 .with_leaves(keys2.into_iter().zip(leaves2.into_iter()))
419 .unwrap();
420
421 let idx = NodeIndex::new(1, 0);
422 assert_eq!(
423 store.get_node(smt.root(), idx).unwrap(),
424 smt.get_node(&idx).unwrap()
425 );
426
427 Ok(())
428 }
429
430 #[test]
431 fn test_add_sparse_merkle_tree() -> Result<(), MerkleError> {
432 let mut store = MerkleStore::default();
433 store.add_sparse_merkle_tree(KEYS4.into_iter().zip(LEAVES4.into_iter()))?;
434
435 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)
436 .unwrap()
437 .with_leaves(KEYS4.into_iter().zip(LEAVES4.into_iter()))
438 .unwrap();
439
440 let idx = NodeIndex::new(1, 0);
441 assert_eq!(
442 store.get_node(smt.root(), idx).unwrap(),
443 smt.get_node(&idx).unwrap()
444 );
445 let idx = NodeIndex::new(1, 1);
446 assert_eq!(
447 store.get_node(smt.root(), idx).unwrap(),
448 smt.get_node(&idx).unwrap()
449 );
450
451 Ok(())
452 }
453
454 #[test]
455 fn test_add_merkle_paths() -> Result<(), MerkleError> {
456 let mut store = MerkleStore::default();
457 let mtree = MerkleTree::new(LEAVES4.to_vec())?;
458
459 let i0 = 0;
460 let p0 = mtree.get_path(NodeIndex::new(2, i0)).unwrap();
461
462 let i1 = 1;
463 let p1 = mtree.get_path(NodeIndex::new(2, i1)).unwrap();
464
465 let i2 = 2;
466 let p2 = mtree.get_path(NodeIndex::new(2, i2)).unwrap();
467
468 let i3 = 3;
469 let p3 = mtree.get_path(NodeIndex::new(2, i3)).unwrap();
470
471 let paths = [
472 (i0, LEAVES4[i0 as usize], p0),
473 (i1, LEAVES4[i1 as usize], p1),
474 (i2, LEAVES4[i2 as usize], p2),
475 (i3, LEAVES4[i3 as usize], p3),
476 ];
477
478 store
479 .add_merkle_paths(paths.clone())
480 .expect("the valid paths must work");
481
482 let set = MerklePathSet::new(3).with_paths(paths).unwrap();
483
484 assert_eq!(
485 set.get_node(NodeIndex::new(3, 0)).unwrap(),
486 store.get_node(set.root(), NodeIndex::new(2, 0b00)).unwrap(),
487 );
488 assert_eq!(
489 set.get_node(NodeIndex::new(3, 1)).unwrap(),
490 store.get_node(set.root(), NodeIndex::new(2, 0b01)).unwrap(),
491 );
492 assert_eq!(
493 set.get_node(NodeIndex::new(3, 2)).unwrap(),
494 store.get_node(set.root(), NodeIndex::new(2, 0b10)).unwrap(),
495 );
496 assert_eq!(
497 set.get_node(NodeIndex::new(3, 3)).unwrap(),
498 store.get_node(set.root(), NodeIndex::new(2, 0b11)).unwrap(),
499 );
500
501 Ok(())
502 }
503}