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