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