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