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