miden-crypto/src/merkle/store/mod.rs
src/merkle/store/mod.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, MerklePathSet, MerkleTree,
8 NodeIndex, RootPath, Rpo256, RpoDigest, SimpleSmt, ValuePath, Vec, Word,
9};
10
11#[cfg(test)]
12mod tests;
13
14#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
15pub struct Node {
16 left: RpoDigest,
17 right: RpoDigest,
18}
19
20#[derive(Debug, Clone, Eq, PartialEq)]
21pub struct MerkleStore {
22 nodes: BTreeMap<RpoDigest, Node>,
23}
24
25impl Default for MerkleStore {
26 fn default() -> Self {
27 Self::new()
28 }
29}
30
31impl MerkleStore {
32 // CONSTRUCTORS
33 // --------------------------------------------------------------------------------------------
34
35 /// Creates an empty `MerkleStore` instance.
36 pub fn new() -> MerkleStore {
37 // pre-populate the store with the empty hashes
38 let subtrees = EmptySubtreeRoots::empty_hashes(64);
39 let nodes = subtrees
40 .iter()
41 .copied()
42 .zip(subtrees.iter().skip(1).copied())
43 .map(|(child, parent)| {
44 (
45 parent,
46 Node {
47 left: child,
48 right: child,
49 },
50 )
51 })
52 .collect();
53
54 MerkleStore { nodes }
55 }
56
57 /// Appends the provided merkle tree represented by its `leaves` to the set.
58 pub fn with_merkle_tree<I>(mut self, leaves: I) -> Result<Self, MerkleError>
59 where
60 I: IntoIterator<Item = Word>,
61 {
62 self.add_merkle_tree(leaves)?;
63 Ok(self)
64 }
65
66 /// Appends the provided sparse merkle tree represented by its `entries` to the set.
67 pub fn with_sparse_merkle_tree<R, I>(mut self, entries: R) -> Result<Self, MerkleError>
68 where
69 R: IntoIterator<IntoIter = I>,
70 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
71 {
72 self.add_sparse_merkle_tree(entries)?;
73 Ok(self)
74 }
75
76 /// Appends the provided merkle path set.
77 pub fn with_merkle_path(
78 mut self,
79 index_value: u64,
80 node: Word,
81 path: MerklePath,
82 ) -> Result<Self, MerkleError> {
83 self.add_merkle_path(index_value, node, path)?;
84 Ok(self)
85 }
86
87 /// Appends the provided merkle path set.
88 pub fn with_merkle_paths<I>(mut self, paths: I) -> Result<Self, MerkleError>
89 where
90 I: IntoIterator<Item = (u64, Word, MerklePath)>,
91 {
92 self.add_merkle_paths(paths)?;
93 Ok(self)
94 }
95
96 // PUBLIC ACCESSORS
97 // --------------------------------------------------------------------------------------------
98
99 /// Returns the node at `index` rooted on the tree `root`.
100 ///
101 /// # Errors
102 ///
103 /// This method can return the following errors:
104 /// - `RootNotInStore` if the `root` is not present in the store.
105 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
106 pub fn get_node(&self, root: Word, index: NodeIndex) -> Result<Word, MerkleError> {
107 let mut hash: RpoDigest = root.into();
108
109 // corner case: check the root is in the store when called with index `NodeIndex::root()`
110 self.nodes
111 .get(&hash)
112 .ok_or(MerkleError::RootNotInStore(hash.into()))?;
113
114 for bit in index.bit_iterator().rev() {
115 let node = self
116 .nodes
117 .get(&hash)
118 .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
119 hash = if bit { node.right } else { node.left }
120 }
121
122 Ok(hash.into())
123 }
124
125 /// Returns the node at the specified `index` and its opening to the `root`.
126 ///
127 /// The path starts at the sibling of the target leaf.
128 ///
129 /// # Errors
130 ///
131 /// This method can return the following errors:
132 /// - `RootNotInStore` if the `root` is not present in the store.
133 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
134 pub fn get_path(&self, root: Word, index: NodeIndex) -> Result<ValuePath, MerkleError> {
135 let mut hash: RpoDigest = root.into();
136 let mut path = Vec::with_capacity(index.depth().into());
137
138 // corner case: check the root is in the store when called with index `NodeIndex::root()`
139 self.nodes
140 .get(&hash)
141 .ok_or(MerkleError::RootNotInStore(hash.into()))?;
142
143 for bit in index.bit_iterator().rev() {
144 let node = self
145 .nodes
146 .get(&hash)
147 .ok_or(MerkleError::NodeNotInStore(hash.into(), index))?;
148
149 hash = if bit {
150 path.push(node.left.into());
151 node.right
152 } else {
153 path.push(node.right.into());
154 node.left
155 }
156 }
157
158 // the path is computed from root to leaf, so it must be reversed
159 path.reverse();
160
161 Ok(ValuePath {
162 value: hash.into(),
163 path: MerklePath::new(path),
164 })
165 }
166
167 // STATE MUTATORS
168 // --------------------------------------------------------------------------------------------
169
170 /// Adds all the nodes of a Merkle tree represented by `leaves`.
171 ///
172 /// This will instantiate a Merkle tree using `leaves` and include all the nodes into the
173 /// store.
174 ///
175 /// # Errors
176 ///
177 /// This method may return the following errors:
178 /// - `DepthTooSmall` if leaves is empty or contains only 1 element
179 /// - `NumLeavesNotPowerOfTwo` if the number of leaves is not a power-of-two
180 pub fn add_merkle_tree<I>(&mut self, leaves: I) -> Result<Word, MerkleError>
181 where
182 I: IntoIterator<Item = Word>,
183 {
184 let leaves: Vec<_> = leaves.into_iter().collect();
185 if leaves.len() < 2 {
186 return Err(MerkleError::DepthTooSmall(leaves.len() as u8));
187 }
188
189 let layers = leaves.len().ilog2();
190 let tree = MerkleTree::new(leaves)?;
191
192 let mut depth = 0;
193 let mut parent_offset = 1;
194 let mut child_offset = 2;
195 while depth < layers {
196 let layer_size = 1usize << depth;
197 for _ in 0..layer_size {
198 // merkle tree is using level form representation, so left and right siblings are
199 // next to each other
200 let left = tree.nodes[child_offset];
201 let right = tree.nodes[child_offset + 1];
202 self.nodes.insert(
203 tree.nodes[parent_offset].into(),
204 Node {
205 left: left.into(),
206 right: right.into(),
207 },
208 );
209 parent_offset += 1;
210 child_offset += 2;
211 }
212 depth += 1;
213 }
214
215 Ok(tree.nodes[1])
216 }
217
218 /// Adds all the nodes of a Sparse Merkle tree represented by `entries`.
219 ///
220 /// This will instantiate a Sparse Merkle tree using `entries` and include all the nodes into
221 /// the store.
222 ///
223 /// # Errors
224 ///
225 /// This will return `InvalidEntriesCount` if the length of `entries` is not `63`.
226 pub fn add_sparse_merkle_tree<R, I>(&mut self, entries: R) -> Result<Word, MerkleError>
227 where
228 R: IntoIterator<IntoIter = I>,
229 I: Iterator<Item = (u64, Word)> + ExactSizeIterator,
230 {
231 let smt = SimpleSmt::new(SimpleSmt::MAX_DEPTH)?.with_leaves(entries)?;
232 for branch in smt.store.branches.values() {
233 let parent = Rpo256::merge(&[branch.left, branch.right]);
234 self.nodes.insert(
235 parent,
236 Node {
237 left: branch.left,
238 right: branch.right,
239 },
240 );
241 }
242
243 Ok(smt.root())
244 }
245
246 /// Adds all the nodes of a Merkle path represented by `path`, opening to `node`. Returns the
247 /// new root.
248 ///
249 /// This will compute the sibling elements determined by the Merkle `path` and `node`, and
250 /// include all the nodes into the store.
251 pub fn add_merkle_path(
252 &mut self,
253 index_value: u64,
254 mut node: Word,
255 path: MerklePath,
256 ) -> Result<Word, MerkleError> {
257 let mut index = NodeIndex::new(self.nodes.len() as u8, index_value);
258
259 for sibling in path {
260 let (left, right) = match index.is_value_odd() {
261 true => (sibling, node),
262 false => (node, sibling),
263 };
264 let parent = Rpo256::merge(&[left.into(), right.into()]);
265 self.nodes.insert(
266 parent,
267 Node {
268 left: left.into(),
269 right: right.into(),
270 },
271 );
272
273 index.move_up();
274 node = parent.into();
275 }
276
277 Ok(node)
278 }
279
280 /// Adds all the nodes of multiple Merkle paths into the store.
281 ///
282 /// This will compute the sibling elements for each Merkle `path` and include all the nodes
283 /// into the store.
284 ///
285 /// For further reference, check [MerkleStore::add_merkle_path].
286 ///
287 /// # Errors
288 ///
289 /// Every path must resolve to the same root, otherwise this will return an `ConflictingRoots`
290 /// error.
291 pub fn add_merkle_paths<I>(&mut self, paths: I) -> Result<Word, MerkleError>
292 where
293 I: IntoIterator<Item = (u64, Word, MerklePath)>,
294 {
295 let paths: Vec<(u64, Word, MerklePath)> = paths.into_iter().collect();
296
297 let roots: BTreeSet<RpoDigest> = paths
298 .iter()
299 .map(|(index, node, path)| path.compute_root(*index, *node).into())
300 .collect();
301
302 if roots.len() != 1 {
303 return Err(MerkleError::ConflictingRoots(
304 roots.iter().map(|v| Word::from(*v)).collect(),
305 ));
306 }
307
308 for (index_value, node, path) in paths {
309 self.add_merkle_path(index_value, node, path)?;
310 }
311
312 Ok(roots.iter().next().unwrap().into())
313 }
314
315 /// Appends the provided [MerklePathSet] into the store.
316 ///
317 /// For further reference, check [MerkleStore::add_merkle_path].
318 pub fn add_merkle_path_set(&mut self, path_set: &MerklePathSet) -> Result<Word, MerkleError> {
319 let root = path_set.root();
320 for (index, path) in path_set.to_paths() {
321 self.add_merkle_path(index, path.value, path.path)?;
322 }
323 Ok(root)
324 }
325
326 /// Sets a node to `value`.
327 ///
328 /// # Errors
329 ///
330 /// This method can return the following errors:
331 /// - `RootNotInStore` if the `root` is not present in the store.
332 /// - `NodeNotInStore` if a node needed to traverse from `root` to `index` is not present in the store.
333 pub fn set_node(
334 &mut self,
335 mut root: Word,
336 index: NodeIndex,
337 value: Word,
338 ) -> Result<RootPath, MerkleError> {
339 let node = value;
340 let ValuePath { value, path } = self.get_path(root, index)?;
341
342 // performs the update only if the node value differs from the opening
343 if node != value {
344 root = self.add_merkle_path(index.value(), node, path.clone())?;
345 }
346
347 Ok(RootPath { root, path })
348 }
349
350 pub fn merge_roots(&mut self, root1: Word, root2: Word) -> Result<Word, MerkleError> {
351 let root1: RpoDigest = root1.into();
352 let root2: RpoDigest = root2.into();
353
354 if !self.nodes.contains_key(&root1) {
355 Err(MerkleError::NodeNotInStore(
356 root1.into(),
357 NodeIndex::new(0, 0),
358 ))
359 } else if !self.nodes.contains_key(&root1) {
360 Err(MerkleError::NodeNotInStore(
361 root2.into(),
362 NodeIndex::new(0, 0),
363 ))
364 } else {
365 let parent: Word = Rpo256::merge(&[root1, root2]).into();
366 self.nodes.insert(
367 parent.into(),
368 Node {
369 left: root1,
370 right: root2,
371 },
372 );
373
374 Ok(parent)
375 }
376 }
377}