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