miden_crypto/merkle/
sparse_path.rs

1use alloc::{borrow::Cow, vec::Vec};
2use core::{
3    iter::{self, FusedIterator},
4    num::NonZero,
5};
6
7use super::{
8    EmptySubtreeRoots, InnerNodeInfo, MerkleError, MerklePath, NodeIndex, Word, smt::SMT_MAX_DEPTH,
9};
10use crate::{
11    hash::rpo::Rpo256,
12    utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
13};
14
15/// A different representation of [`MerklePath`] designed for memory efficiency for Merkle paths
16/// with empty nodes.
17///
18/// Empty nodes in the path are stored only as their position, represented with a bitmask. A
19/// maximum of 64 nodes (`SMT_MAX_DEPTH`) can be stored (empty and non-empty). The more nodes in a
20/// path are empty, the less memory this struct will use. This type calculates empty nodes on-demand
21/// when iterated through, converted to a [MerklePath], or an empty node is retrieved with
22/// [`SparseMerklePath::at_depth()`], which will incur overhead.
23///
24/// NOTE: This type assumes that Merkle paths always span from the root of the tree to a leaf.
25/// Partial paths are not supported.
26#[derive(Clone, Debug, Default, PartialEq, Eq)]
27#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
28pub struct SparseMerklePath {
29    /// A bitmask representing empty nodes. The set bit corresponds to the depth of an empty node.
30    /// The least significant bit (bit 0) describes depth 1 node (root's children).
31    /// The `bit index + 1` is equal to node's depth.
32    empty_nodes_mask: u64,
33    /// The non-empty nodes, stored in depth-order, but not contiguous across depth.
34    nodes: Vec<Word>,
35}
36
37impl SparseMerklePath {
38    /// Constructs a new sparse Merkle path from a bitmask of empty nodes and a vector of non-empty
39    /// nodes.
40    ///
41    /// The `empty_nodes_mask` is a bitmask where each set bit indicates that the node at that
42    /// depth is empty. The least significant bit (bit 0) describes depth 1 node (root's children).
43    /// The `bit index + 1` is equal to node's depth.
44    /// The `nodes` vector must contain the non-empty nodes in depth order.
45    ///
46    /// # Errors
47    /// - [MerkleError::InvalidPathLength] if the provided `nodes` vector is shorter than the
48    ///   minimum length required by the `empty_nodes_mask`.
49    /// - [MerkleError::DepthTooBig] if the total depth of the path (calculated from the
50    ///   `empty_nodes_mask` and `nodes`) is greater than [SMT_MAX_DEPTH].
51    pub fn from_parts(empty_nodes_mask: u64, nodes: Vec<Word>) -> Result<Self, MerkleError> {
52        // The most significant set bit in the mask marks the minimum length of the path.
53        // For every zero bit before the first set bit, there must be a corresponding node in
54        // `nodes`.
55        // For example, if the mask is `0b1100`, this means that the first two nodes
56        // (depths 1 and 2) are non-empty, and the next two nodes (depths 3 and 4) are empty.
57        // The minimum length of the path is 4, and the `nodes` vector must contain at least 2
58        // nodes to account for the first two zeroes in the mask (depths 1 and 2).
59        let min_path_len = u64::BITS - empty_nodes_mask.leading_zeros();
60        let empty_nodes_count = empty_nodes_mask.count_ones();
61        let min_non_empty_nodes = (min_path_len - empty_nodes_count) as usize;
62
63        if nodes.len() < min_non_empty_nodes {
64            return Err(MerkleError::InvalidPathLength(min_non_empty_nodes));
65        }
66
67        let depth = Self::depth_from_parts(empty_nodes_mask, &nodes) as u8;
68        if depth > SMT_MAX_DEPTH {
69            return Err(MerkleError::DepthTooBig(depth as u64));
70        }
71
72        Ok(Self { empty_nodes_mask, nodes })
73    }
74
75    /// Constructs a sparse Merkle path from an iterator over Merkle nodes that also knows its
76    /// exact size (such as iterators created with [Vec::into_iter]). The iterator must be in order
77    /// of deepest to shallowest.
78    ///
79    /// Knowing the size is necessary to calculate the depth of the tree, which is needed to detect
80    /// which nodes are empty nodes.
81    ///
82    /// # Errors
83    /// Returns [MerkleError::DepthTooBig] if `tree_depth` is greater than [SMT_MAX_DEPTH].
84    pub fn from_sized_iter<I>(iterator: I) -> Result<Self, MerkleError>
85    where
86        I: IntoIterator<IntoIter: ExactSizeIterator, Item = Word>,
87    {
88        let iterator = iterator.into_iter();
89        let tree_depth = iterator.len() as u8;
90
91        if tree_depth > SMT_MAX_DEPTH {
92            return Err(MerkleError::DepthTooBig(tree_depth as u64));
93        }
94
95        let mut empty_nodes_mask: u64 = 0;
96        let mut nodes: Vec<Word> = Default::default();
97
98        for (depth, node) in iter::zip(path_depth_iter(tree_depth), iterator) {
99            let &equivalent_empty_node = EmptySubtreeRoots::entry(tree_depth, depth.get());
100            let is_empty = node == equivalent_empty_node;
101            let node = if is_empty { None } else { Some(node) };
102
103            match node {
104                Some(node) => nodes.push(node),
105                None => empty_nodes_mask |= Self::bitmask_for_depth(depth),
106            }
107        }
108
109        Ok(SparseMerklePath { nodes, empty_nodes_mask })
110    }
111
112    /// Returns the total depth of this path, i.e., the number of nodes this path represents.
113    pub fn depth(&self) -> u8 {
114        Self::depth_from_parts(self.empty_nodes_mask, &self.nodes) as u8
115    }
116
117    /// Get a specific node in this path at a given depth.
118    ///
119    /// The `depth` parameter is defined in terms of `self.depth()`. Merkle paths conventionally do
120    /// not include the root, so the shallowest depth is `1`, and the deepest depth is
121    /// `self.depth()`.
122    ///
123    /// # Errors
124    /// Returns [MerkleError::DepthTooBig] if `node_depth` is greater than the total depth of this
125    /// path.
126    pub fn at_depth(&self, node_depth: NonZero<u8>) -> Result<Word, MerkleError> {
127        if node_depth.get() > self.depth() {
128            return Err(MerkleError::DepthTooBig(node_depth.get().into()));
129        }
130
131        let node = if let Some(nonempty_index) = self.get_nonempty_index(node_depth) {
132            self.nodes[nonempty_index]
133        } else {
134            *EmptySubtreeRoots::entry(self.depth(), node_depth.get())
135        };
136
137        Ok(node)
138    }
139
140    /// Deconstructs this path into its component parts.
141    ///
142    /// Returns a tuple containing:
143    /// - a bitmask where each set bit indicates that the node at that depth is empty. The least
144    ///   significant bit (bit 0) describes depth 1 node (root's children).
145    /// - a vector of non-empty nodes in depth order.
146    pub fn into_parts(self) -> (u64, Vec<Word>) {
147        (self.empty_nodes_mask, self.nodes)
148    }
149
150    // PROVIDERS
151    // ============================================================================================
152
153    /// Constructs a borrowing iterator over the nodes in this path.
154    /// Starts from the leaf and iterates toward the root (excluding the root).
155    pub fn iter(&self) -> impl ExactSizeIterator<Item = Word> {
156        self.into_iter()
157    }
158
159    /// Computes the Merkle root for this opening.
160    pub fn compute_root(&self, index: u64, node_to_prove: Word) -> Result<Word, MerkleError> {
161        let mut index = NodeIndex::new(self.depth(), index)?;
162        let root = self.iter().fold(node_to_prove, |node, sibling| {
163            // Compute the node and move to the next iteration.
164            let children = index.build_node(node, sibling);
165            index.move_up();
166            Rpo256::merge(&children)
167        });
168
169        Ok(root)
170    }
171
172    /// Verifies the Merkle opening proof towards the provided root.
173    ///
174    /// # Errors
175    /// Returns an error if:
176    /// - provided node index is invalid.
177    /// - root calculated during the verification differs from the provided one.
178    pub fn verify(&self, index: u64, node: Word, &expected_root: &Word) -> Result<(), MerkleError> {
179        let computed_root = self.compute_root(index, node)?;
180        if computed_root != expected_root {
181            return Err(MerkleError::ConflictingRoots {
182                expected_root,
183                actual_root: computed_root,
184            });
185        }
186
187        Ok(())
188    }
189
190    /// Given the node this path opens to, return an iterator of all the nodes that are known via
191    /// this path.
192    ///
193    /// Each item in the iterator is an [InnerNodeInfo], containing the hash of a node as `.value`,
194    /// and its two children as `.left` and `.right`. The very first item in that iterator will be
195    /// the parent of `node_to_prove` as stored in this [SparseMerklePath].
196    ///
197    /// From there, the iterator will continue to yield every further parent and both of its
198    /// children, up to and including the root node.
199    ///
200    /// If `node_to_prove` is not the node this path is an opening to, or `index` is not the
201    /// correct index for that node, the returned nodes will be meaningless.
202    ///
203    /// # Errors
204    /// Returns an error if the specified index is not valid for this path.
205    pub fn authenticated_nodes(
206        &self,
207        index: u64,
208        node_to_prove: Word,
209    ) -> Result<InnerNodeIterator<'_>, MerkleError> {
210        let index = NodeIndex::new(self.depth(), index)?;
211        Ok(InnerNodeIterator { path: self, index, value: node_to_prove })
212    }
213
214    // PRIVATE HELPERS
215    // ============================================================================================
216
217    const fn bitmask_for_depth(node_depth: NonZero<u8>) -> u64 {
218        // - 1 because paths do not include the root.
219        1 << (node_depth.get() - 1)
220    }
221
222    const fn is_depth_empty(&self, node_depth: NonZero<u8>) -> bool {
223        (self.empty_nodes_mask & Self::bitmask_for_depth(node_depth)) != 0
224    }
225
226    /// Index of the non-empty node in the `self.nodes` vector. If the specified depth is
227    /// empty, None is returned.
228    fn get_nonempty_index(&self, node_depth: NonZero<u8>) -> Option<usize> {
229        if self.is_depth_empty(node_depth) {
230            return None;
231        }
232
233        let bit_index = node_depth.get() - 1;
234        let without_shallower = self.empty_nodes_mask >> bit_index;
235        let empty_deeper = without_shallower.count_ones() as usize;
236        // The vec index we would use if we didn't have any empty nodes to account for...
237        let normal_index = (self.depth() - node_depth.get()) as usize;
238        // subtracted by the number of empty nodes that are deeper than us.
239        Some(normal_index - empty_deeper)
240    }
241
242    /// Returns the total depth of this path from its parts.
243    fn depth_from_parts(empty_nodes_mask: u64, nodes: &[Word]) -> usize {
244        nodes.len() + empty_nodes_mask.count_ones() as usize
245    }
246}
247
248// SERIALIZATION
249// ================================================================================================
250
251impl Serializable for SparseMerklePath {
252    fn write_into<W: ByteWriter>(&self, target: &mut W) {
253        target.write_u8(self.depth());
254        target.write_u64(self.empty_nodes_mask);
255        target.write_many(&self.nodes);
256    }
257}
258
259impl Deserializable for SparseMerklePath {
260    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
261        let depth = source.read_u8()?;
262        if depth > SMT_MAX_DEPTH {
263            return Err(DeserializationError::InvalidValue(format!(
264                "SparseMerklePath max depth exceeded ({depth} > {SMT_MAX_DEPTH})",
265            )));
266        }
267        let empty_nodes_mask = source.read_u64()?;
268        let empty_nodes_count = empty_nodes_mask.count_ones();
269        if empty_nodes_count > depth as u32 {
270            return Err(DeserializationError::InvalidValue(format!(
271                "SparseMerklePath has more empty nodes ({empty_nodes_count}) than its full length ({depth})",
272            )));
273        }
274        let count = depth as u32 - empty_nodes_count;
275        let nodes: Vec<Word> = source.read_many_iter(count as usize)?.collect::<Result<_, _>>()?;
276        Ok(Self { empty_nodes_mask, nodes })
277    }
278}
279
280// CONVERSIONS
281// ================================================================================================
282
283impl From<SparseMerklePath> for MerklePath {
284    fn from(sparse_path: SparseMerklePath) -> Self {
285        MerklePath::from_iter(sparse_path)
286    }
287}
288
289impl TryFrom<MerklePath> for SparseMerklePath {
290    type Error = MerkleError;
291
292    /// # Errors
293    ///
294    /// This conversion returns [MerkleError::DepthTooBig] if the path length is greater than
295    /// [`SMT_MAX_DEPTH`].
296    fn try_from(path: MerklePath) -> Result<Self, MerkleError> {
297        SparseMerklePath::from_sized_iter(path)
298    }
299}
300
301impl From<SparseMerklePath> for Vec<Word> {
302    fn from(path: SparseMerklePath) -> Self {
303        Vec::from_iter(path)
304    }
305}
306
307// ITERATORS
308// ================================================================================================
309
310/// Iterator for [`SparseMerklePath`]. Starts from the leaf and iterates toward the root (excluding
311/// the root).
312pub struct SparseMerklePathIter<'p> {
313    /// The "inner" value we're iterating over.
314    path: Cow<'p, SparseMerklePath>,
315
316    /// The depth a `next()` call will get. `next_depth == 0` indicates that the iterator has been
317    /// exhausted.
318    next_depth: u8,
319}
320
321impl Iterator for SparseMerklePathIter<'_> {
322    type Item = Word;
323
324    fn next(&mut self) -> Option<Word> {
325        let this_depth = self.next_depth;
326        // Paths don't include the root, so if `this_depth` is 0 then we keep returning `None`.
327        let this_depth = NonZero::new(this_depth)?;
328        self.next_depth = this_depth.get() - 1;
329
330        // `this_depth` is only ever decreasing, so it can't ever exceed `self.path.depth()`.
331        let node = self
332            .path
333            .at_depth(this_depth)
334            .expect("current depth should never exceed the path depth");
335        Some(node)
336    }
337
338    // SparseMerkleIter always knows its exact size.
339    fn size_hint(&self) -> (usize, Option<usize>) {
340        let remaining = ExactSizeIterator::len(self);
341        (remaining, Some(remaining))
342    }
343}
344
345impl ExactSizeIterator for SparseMerklePathIter<'_> {
346    fn len(&self) -> usize {
347        self.next_depth as usize
348    }
349}
350
351impl FusedIterator for SparseMerklePathIter<'_> {}
352
353// TODO: impl DoubleEndedIterator.
354
355impl IntoIterator for SparseMerklePath {
356    type IntoIter = SparseMerklePathIter<'static>;
357    type Item = <Self::IntoIter as Iterator>::Item;
358
359    fn into_iter(self) -> SparseMerklePathIter<'static> {
360        let tree_depth = self.depth();
361        SparseMerklePathIter {
362            path: Cow::Owned(self),
363            next_depth: tree_depth,
364        }
365    }
366}
367
368impl<'p> IntoIterator for &'p SparseMerklePath {
369    type Item = <SparseMerklePathIter<'p> as Iterator>::Item;
370    type IntoIter = SparseMerklePathIter<'p>;
371
372    fn into_iter(self) -> SparseMerklePathIter<'p> {
373        let tree_depth = self.depth();
374        SparseMerklePathIter {
375            path: Cow::Borrowed(self),
376            next_depth: tree_depth,
377        }
378    }
379}
380
381/// An iterator over nodes known by a [SparseMerklePath]. See
382/// [`SparseMerklePath::authenticated_nodes()`].
383pub struct InnerNodeIterator<'p> {
384    path: &'p SparseMerklePath,
385    index: NodeIndex,
386    value: Word,
387}
388
389impl Iterator for InnerNodeIterator<'_> {
390    type Item = InnerNodeInfo;
391
392    fn next(&mut self) -> Option<Self::Item> {
393        if self.index.is_root() {
394            return None;
395        }
396
397        let index_depth = NonZero::new(self.index.depth()).expect("non-root depth cannot be 0");
398        let path_node = self.path.at_depth(index_depth).unwrap();
399
400        let children = self.index.build_node(self.value, path_node);
401        self.value = Rpo256::merge(&children);
402        self.index.move_up();
403
404        Some(InnerNodeInfo {
405            value: self.value,
406            left: children[0],
407            right: children[1],
408        })
409    }
410}
411
412// COMPARISONS
413// ================================================================================================
414impl PartialEq<MerklePath> for SparseMerklePath {
415    fn eq(&self, rhs: &MerklePath) -> bool {
416        if self.depth() != rhs.depth() {
417            return false;
418        }
419
420        for (node, &rhs_node) in iter::zip(self, rhs.iter()) {
421            if node != rhs_node {
422                return false;
423            }
424        }
425
426        true
427    }
428}
429
430impl PartialEq<SparseMerklePath> for MerklePath {
431    fn eq(&self, rhs: &SparseMerklePath) -> bool {
432        rhs == self
433    }
434}
435
436// HELPERS
437// ================================================================================================
438
439/// Iterator for path depths, which start at the deepest part of the tree and go the shallowest
440/// depth before the root (depth 1).
441fn path_depth_iter(tree_depth: u8) -> impl ExactSizeIterator<Item = NonZero<u8>> {
442    let top_down_iter = (1..=tree_depth).map(|depth| {
443        // SAFETY: `RangeInclusive<1, _>` cannot ever yield 0. Even if `tree_depth` is 0, then the
444        // range is `RangeInclusive<1, 0>` will simply not yield any values, and this block won't
445        // even be reached.
446        unsafe { NonZero::new_unchecked(depth) }
447    });
448
449    // Reverse the top-down iterator to get a bottom-up iterator.
450    top_down_iter.rev()
451}
452
453// TESTS
454// ================================================================================================
455#[cfg(test)]
456mod tests {
457    use alloc::vec::Vec;
458    use core::num::NonZero;
459
460    use assert_matches::assert_matches;
461    use p3_field::PrimeCharacteristicRing;
462
463    use super::SparseMerklePath;
464    use crate::{
465        Felt, ONE, Word,
466        merkle::{
467            EmptySubtreeRoots, MerkleError, MerklePath, MerkleTree, NodeIndex,
468            smt::{LeafIndex, SMT_MAX_DEPTH, SimpleSmt, Smt, SparseMerkleTree},
469            sparse_path::path_depth_iter,
470        },
471    };
472
473    fn make_smt(pair_count: u64) -> Smt {
474        let entries: Vec<(Word, Word)> = (0..pair_count)
475            .map(|n| {
476                let leaf_index = ((n as f64 / pair_count as f64) * 255.0) as u64;
477                let key = Word::new([ONE, ONE, Felt::new(n), Felt::new(leaf_index)]);
478                let value = Word::new([ONE, ONE, ONE, ONE]);
479                (key, value)
480            })
481            .collect();
482
483        Smt::with_entries(entries).unwrap()
484    }
485
486    /// Manually test the exact bit patterns for a sample path of 8 nodes, including both empty and
487    /// non-empty nodes.
488    ///
489    /// This also offers an overview of what each part of the bit-math involved means and
490    /// represents.
491    #[test]
492    fn test_sparse_bits() {
493        const DEPTH: u8 = 8;
494        let raw_nodes: [Word; DEPTH as usize] = [
495            // Depth 8.
496            ([8u8, 8, 8, 8].into()),
497            // Depth 7.
498            *EmptySubtreeRoots::entry(DEPTH, 7),
499            // Depth 6.
500            *EmptySubtreeRoots::entry(DEPTH, 6),
501            // Depth 5.
502            [5u8, 5, 5, 5].into(),
503            // Depth 4.
504            [4u8, 4, 4, 4].into(),
505            // Depth 3.
506            *EmptySubtreeRoots::entry(DEPTH, 3),
507            // Depth 2.
508            *EmptySubtreeRoots::entry(DEPTH, 2),
509            // Depth 1.
510            *EmptySubtreeRoots::entry(DEPTH, 1),
511            // Root is not included.
512        ];
513
514        let sparse_nodes: [Option<Word>; DEPTH as usize] = [
515            // Depth 8.
516            Some([8u8, 8, 8, 8].into()),
517            // Depth 7.
518            None,
519            // Depth 6.
520            None,
521            // Depth 5.
522            Some([5u8, 5, 5, 5].into()),
523            // Depth 4.
524            Some([4u8, 4, 4, 4].into()),
525            // Depth 3.
526            None,
527            // Depth 2.
528            None,
529            // Depth 1.
530            None,
531            // Root is not included.
532        ];
533
534        const EMPTY_BITS: u64 = 0b0110_0111;
535
536        let sparse_path = SparseMerklePath::from_sized_iter(raw_nodes).unwrap();
537
538        assert_eq!(sparse_path.empty_nodes_mask, EMPTY_BITS);
539
540        // Keep track of how many non-empty nodes we have seen
541        let mut nonempty_idx = 0;
542
543        // Test starting from the deepest nodes (depth 8)
544        for depth in (1..=8).rev() {
545            let idx = (sparse_path.depth() - depth) as usize;
546            let bit = 1 << (depth - 1);
547
548            // Check that the depth bit is set correctly...
549            let is_set = (sparse_path.empty_nodes_mask & bit) != 0;
550            assert_eq!(is_set, sparse_nodes.get(idx).unwrap().is_none());
551
552            if is_set {
553                // Check that we don't return digests for empty nodes
554                let &test_node = sparse_nodes.get(idx).unwrap();
555                assert_eq!(test_node, None);
556            } else {
557                // Check that we can calculate non-empty indices correctly.
558                let control_node = raw_nodes.get(idx).unwrap();
559                assert_eq!(
560                    sparse_path.get_nonempty_index(NonZero::new(depth).unwrap()).unwrap(),
561                    nonempty_idx
562                );
563                let test_node = sparse_path.nodes.get(nonempty_idx).unwrap();
564                assert_eq!(test_node, control_node);
565
566                nonempty_idx += 1;
567            }
568        }
569    }
570
571    #[test]
572    fn from_parts() {
573        const DEPTH: u8 = 8;
574        let raw_nodes: [Word; DEPTH as usize] = [
575            // Depth 8.
576            ([8u8, 8, 8, 8].into()),
577            // Depth 7.
578            *EmptySubtreeRoots::entry(DEPTH, 7),
579            // Depth 6.
580            *EmptySubtreeRoots::entry(DEPTH, 6),
581            // Depth 5.
582            [5u8, 5, 5, 5].into(),
583            // Depth 4.
584            [4u8, 4, 4, 4].into(),
585            // Depth 3.
586            *EmptySubtreeRoots::entry(DEPTH, 3),
587            // Depth 2.
588            *EmptySubtreeRoots::entry(DEPTH, 2),
589            // Depth 1.
590            *EmptySubtreeRoots::entry(DEPTH, 1),
591            // Root is not included.
592        ];
593
594        let empty_nodes_mask = 0b0110_0111;
595        let nodes = vec![[8u8, 8, 8, 8].into(), [5u8, 5, 5, 5].into(), [4u8, 4, 4, 4].into()];
596        let insufficient_nodes = vec![[4u8, 4, 4, 4].into()];
597
598        let error = SparseMerklePath::from_parts(empty_nodes_mask, insufficient_nodes).unwrap_err();
599        assert_matches!(error, MerkleError::InvalidPathLength(2));
600
601        let iter_sparse_path = SparseMerklePath::from_sized_iter(raw_nodes).unwrap();
602        let sparse_path = SparseMerklePath::from_parts(empty_nodes_mask, nodes).unwrap();
603
604        assert_eq!(sparse_path, iter_sparse_path);
605    }
606
607    #[test]
608    fn from_sized_iter() {
609        let tree = make_smt(8192);
610
611        for (key, _value) in tree.entries() {
612            let index = NodeIndex::from(Smt::key_to_leaf_index(key));
613            let sparse_path = tree.get_path(key);
614            for (sparse_node, proof_idx) in
615                itertools::zip_eq(sparse_path.clone(), index.proof_indices())
616            {
617                let proof_node = tree.get_node_hash(proof_idx);
618                assert_eq!(sparse_node, proof_node);
619            }
620        }
621    }
622
623    #[test]
624    fn test_zero_sized() {
625        let nodes: Vec<Word> = Default::default();
626
627        // Sparse paths that don't actually contain any nodes should still be well behaved.
628        let sparse_path = SparseMerklePath::from_sized_iter(nodes).unwrap();
629        assert_eq!(sparse_path.depth(), 0);
630        assert_matches!(
631            sparse_path.at_depth(NonZero::new(1).unwrap()),
632            Err(MerkleError::DepthTooBig(1))
633        );
634        assert_eq!(sparse_path.iter().next(), None);
635        assert_eq!(sparse_path.into_iter().next(), None);
636    }
637
638    use proptest::prelude::*;
639
640    // Arbitrary instance for MerklePath
641    impl Arbitrary for MerklePath {
642        type Parameters = ();
643        type Strategy = BoxedStrategy<Self>;
644
645        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
646            prop::collection::vec(any::<Word>(), 0..=SMT_MAX_DEPTH as usize)
647                .prop_map(MerklePath::new)
648                .boxed()
649        }
650    }
651
652    // Arbitrary instance for SparseMerklePath
653    impl Arbitrary for SparseMerklePath {
654        type Parameters = ();
655        type Strategy = BoxedStrategy<Self>;
656
657        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
658            (0..=SMT_MAX_DEPTH as usize)
659                .prop_flat_map(|depth| {
660                    // Generate a bitmask for empty nodes - avoid overflow
661                    let max_mask = if depth > 0 && depth < 64 {
662                        (1u64 << depth) - 1
663                    } else if depth == 64 {
664                        u64::MAX
665                    } else {
666                        0
667                    };
668                    let empty_nodes_mask =
669                        prop::num::u64::ANY.prop_map(move |mask| mask & max_mask);
670
671                    // Generate non-empty nodes based on the mask
672                    empty_nodes_mask.prop_flat_map(move |mask| {
673                        let empty_count = mask.count_ones() as usize;
674                        let non_empty_count = depth.saturating_sub(empty_count);
675
676                        prop::collection::vec(any::<Word>(), non_empty_count).prop_map(
677                            move |nodes| SparseMerklePath::from_parts(mask, nodes).unwrap(),
678                        )
679                    })
680                })
681                .boxed()
682        }
683    }
684
685    proptest! {
686        #[test]
687        fn sparse_merkle_path_roundtrip_equivalence(path in any::<MerklePath>()) {
688            // Convert MerklePath to SparseMerklePath and back
689            let sparse_result = SparseMerklePath::try_from(path.clone());
690            if path.depth() <= SMT_MAX_DEPTH {
691                let sparse = sparse_result.unwrap();
692                let reconstructed = MerklePath::from(sparse);
693                prop_assert_eq!(path, reconstructed);
694            } else {
695                prop_assert!(sparse_result.is_err());
696            }
697        }
698    }
699    proptest! {
700
701        #[test]
702        fn merkle_path_roundtrip_equivalence(sparse in any::<SparseMerklePath>()) {
703            // Convert SparseMerklePath to MerklePath and back
704            let merkle = MerklePath::from(sparse.clone());
705            let reconstructed = SparseMerklePath::try_from(merkle.clone()).unwrap();
706            prop_assert_eq!(sparse, reconstructed);
707        }
708    }
709    proptest! {
710
711        #[test]
712        fn path_equivalence_tests(path in any::<MerklePath>(), path2 in any::<MerklePath>()) {
713            if path.depth() > SMT_MAX_DEPTH {
714                return Ok(());
715            }
716
717            let sparse = SparseMerklePath::try_from(path.clone()).unwrap();
718
719            // Depth consistency
720            prop_assert_eq!(path.depth(), sparse.depth());
721
722            // Node access consistency including path_depth_iter
723            if path.depth() > 0 {
724                for depth in path_depth_iter(path.depth()) {
725                    let merkle_node = path.at_depth(depth);
726                    let sparse_node = sparse.at_depth(depth);
727
728                    match (merkle_node, sparse_node) {
729                        (Some(m), Ok(s)) => prop_assert_eq!(m, s),
730                        (None, Err(_)) => {},
731                        _ => prop_assert!(false, "Inconsistent node access at depth {}", depth.get()),
732                    }
733                }
734            }
735
736            // Iterator consistency
737            if path.depth() > 0 {
738                let merkle_nodes: Vec<_> = path.iter().collect();
739                let sparse_nodes: Vec<_> = sparse.iter().collect();
740
741                prop_assert_eq!(merkle_nodes.len(), sparse_nodes.len());
742                for (m, s) in merkle_nodes.iter().zip(sparse_nodes.iter()) {
743                    prop_assert_eq!(*m, s);
744                }
745            }
746
747            // Test equality between different representations
748            if path2.depth() <= SMT_MAX_DEPTH {
749                let sparse2 = SparseMerklePath::try_from(path2.clone()).unwrap();
750                prop_assert_eq!(path == path2, sparse == sparse2);
751                prop_assert_eq!(path == sparse2, sparse == path2);
752            }
753        }
754    }
755    // rather heavy tests
756    proptest! {
757        #![proptest_config(ProptestConfig::with_cases(100))]
758
759        #[test]
760        fn compute_root_consistency(
761            tree_data in any::<RandomMerkleTree>(),
762            node in any::<Word>()
763        ) {
764            let RandomMerkleTree { tree, leaves: _,  indices } = tree_data;
765
766            for &leaf_index in indices.iter() {
767                let path = tree.get_path(NodeIndex::new(tree.depth(), leaf_index).unwrap()).unwrap();
768                let sparse = SparseMerklePath::from_sized_iter(path.clone().into_iter()).unwrap();
769
770                let merkle_root = path.compute_root(leaf_index, node);
771                let sparse_root = sparse.compute_root(leaf_index, node);
772
773                match (merkle_root, sparse_root) {
774                    (Ok(m), Ok(s)) => prop_assert_eq!(m, s),
775                    (Err(e1), Err(e2)) => {
776                        // Both should have the same error type
777                        prop_assert_eq!(format!("{:?}", e1), format!("{:?}", e2));
778                    },
779                    _ => prop_assert!(false, "Inconsistent compute_root results"),
780                }
781            }
782        }
783
784        #[test]
785        fn verify_consistency(
786            tree_data in any::<RandomMerkleTree>(),
787            node in any::<Word>()
788        ) {
789            let RandomMerkleTree { tree, leaves, indices } = tree_data;
790
791            for (i, &leaf_index) in indices.iter().enumerate() {
792                let leaf = leaves[i];
793                let path = tree.get_path(NodeIndex::new(tree.depth(), leaf_index).unwrap()).unwrap();
794                let sparse = SparseMerklePath::from_sized_iter(path.clone().into_iter()).unwrap();
795
796                let root = tree.root();
797
798                let merkle_verify = path.verify(leaf_index, leaf, &root);
799                let sparse_verify = sparse.verify(leaf_index, leaf, &root);
800
801                match (merkle_verify, sparse_verify) {
802                    (Ok(()), Ok(())) => {},
803                    (Err(e1), Err(e2)) => {
804                        // Both should have the same error type
805                        prop_assert_eq!(format!("{:?}", e1), format!("{:?}", e2));
806                    },
807                    _ => prop_assert!(false, "Inconsistent verify results"),
808                }
809
810                // Test with wrong node - both should fail
811                let wrong_verify = path.verify(leaf_index, node, &root);
812                let wrong_sparse_verify = sparse.verify(leaf_index, node, &root);
813
814                match (wrong_verify, wrong_sparse_verify) {
815                    (Ok(()), Ok(())) => prop_assert!(false, "Verification should have failed with wrong node"),
816                    (Err(_), Err(_)) => {},
817                    _ => prop_assert!(false, "Inconsistent verification results with wrong node"),
818                }
819            }
820        }
821
822        #[test]
823        fn authenticated_nodes_consistency(
824            tree_data in any::<RandomMerkleTree>()
825        ) {
826            let RandomMerkleTree { tree, leaves, indices } = tree_data;
827
828            for (i, &leaf_index) in indices.iter().enumerate() {
829                let leaf = leaves[i];
830                let path = tree.get_path(NodeIndex::new(tree.depth(), leaf_index).unwrap()).unwrap();
831                let sparse = SparseMerklePath::from_sized_iter(path.clone().into_iter()).unwrap();
832
833                let merkle_result = path.authenticated_nodes(leaf_index, leaf);
834                let sparse_result = sparse.authenticated_nodes(leaf_index, leaf);
835
836                match (merkle_result, sparse_result) {
837                    (Ok(m_iter), Ok(s_iter)) => {
838                        let merkle_nodes: Vec<_> = m_iter.collect();
839                        let sparse_nodes: Vec<_> = s_iter.collect();
840                        prop_assert_eq!(merkle_nodes.len(), sparse_nodes.len());
841                        for (m, s) in merkle_nodes.iter().zip(sparse_nodes.iter()) {
842                            prop_assert_eq!(m, s);
843                        }
844                    },
845                    (Err(e1), Err(e2)) => {
846                        prop_assert_eq!(format!("{:?}", e1), format!("{:?}", e2));
847                    },
848                    _ => prop_assert!(false, "Inconsistent authenticated_nodes results"),
849                }
850            }
851        }
852    }
853
854    #[test]
855    fn test_api_differences() {
856        // This test documents API differences between MerklePath and SparseMerklePath
857
858        // 1. MerklePath has Deref/DerefMut to Vec<Word> - SparseMerklePath does not
859        let merkle = MerklePath::new(vec![Word::default(); 3]);
860        let _vec_ref: &Vec<Word> = &merkle; // This works due to Deref
861        let _vec_mut: &mut Vec<Word> = &mut merkle.clone(); // This works due to DerefMut
862
863        // 2. SparseMerklePath has from_parts() - MerklePath uses new() or from_iter()
864        let sparse = SparseMerklePath::from_parts(0b101, vec![Word::default(); 2]).unwrap();
865        assert_eq!(sparse.depth(), 4); // depth is 4 because mask has bits set up to depth 4
866
867        // 3. SparseMerklePath has from_sized_iter() - MerklePath uses from_iter()
868        let nodes = vec![Word::default(); 3];
869        let sparse_from_iter = SparseMerklePath::from_sized_iter(nodes.clone()).unwrap();
870        let merkle_from_iter = MerklePath::from_iter(nodes);
871        assert_eq!(sparse_from_iter.depth(), merkle_from_iter.depth());
872    }
873
874    // Arbitrary instance for MerkleTree with random leaves
875    #[derive(Debug, Clone)]
876    struct RandomMerkleTree {
877        tree: MerkleTree,
878        leaves: Vec<Word>,
879        indices: Vec<u64>,
880    }
881
882    impl Arbitrary for RandomMerkleTree {
883        type Parameters = ();
884        type Strategy = BoxedStrategy<Self>;
885
886        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
887            // Generate trees with power-of-2 leaves up to 1024 (2^10)
888            prop::sample::select(&[2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
889                .prop_flat_map(|num_leaves| {
890                    prop::collection::vec(any::<Word>(), num_leaves).prop_map(|leaves| {
891                        let tree = MerkleTree::new(leaves.clone()).unwrap();
892                        let indices: Vec<u64> = (0..leaves.len() as u64).collect();
893                        RandomMerkleTree { tree, leaves, indices }
894                    })
895                })
896                .boxed()
897        }
898    }
899
900    // Arbitrary instance for SimpleSmt with random entries
901    #[derive(Debug, Clone)]
902    struct RandomSimpleSmt {
903        tree: SimpleSmt<10>, // Depth 10 = 1024 leaves
904        entries: Vec<(u64, Word)>,
905    }
906
907    impl Arbitrary for RandomSimpleSmt {
908        type Parameters = ();
909        type Strategy = BoxedStrategy<Self>;
910
911        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
912            (1..=100usize) // 1-100 entries in an 1024-leaf tree
913                .prop_flat_map(|num_entries| {
914                    prop::collection::vec(
915                        (
916                            0..1024u64, // Valid indices for 1024-leaf tree
917                            any::<Word>(),
918                        ),
919                        num_entries,
920                    )
921                    .prop_map(|mut entries| {
922                        // Ensure unique indices to avoid duplicates
923                        let mut seen = alloc::collections::BTreeSet::new();
924                        entries.retain(|(idx, _)| seen.insert(*idx));
925
926                        let mut tree = SimpleSmt::new().unwrap();
927                        for (idx, value) in &entries {
928                            let leaf_idx = LeafIndex::new(*idx).unwrap();
929                            tree.insert(leaf_idx, *value);
930                        }
931                        RandomSimpleSmt { tree, entries }
932                    })
933                })
934                .boxed()
935        }
936    }
937
938    // Arbitrary instance for Smt with random entries
939    #[derive(Debug, Clone)]
940    struct RandomSmt {
941        tree: Smt,
942        entries: Vec<(Word, Word)>,
943    }
944
945    impl Arbitrary for RandomSmt {
946        type Parameters = ();
947        type Strategy = BoxedStrategy<Self>;
948
949        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
950            (1..=100usize) // 1-100 entries in a sparse tree
951                .prop_flat_map(|num_entries| {
952                    prop::collection::vec((any::<u64>(), any::<Word>()), num_entries).prop_map(
953                        |indices_n_values| {
954                            let entries: Vec<(Word, Word)> = indices_n_values
955                                .into_iter()
956                                .enumerate()
957                                .map(|(n, (leaf_index, value))| {
958                                    // SMT uses the most significant element (index 3) as leaf index
959                                    // Ensure we use valid leaf indices for the SMT depth
960                                    let valid_leaf_index = leaf_index % (1u64 << 60); // Use large but valid range
961                                    let key = Word::new([
962                                        Felt::new(n as u64),         // element 0
963                                        Felt::new(n as u64 + 1),     // element 1
964                                        Felt::new(n as u64 + 2),     // element 2
965                                        Felt::new(valid_leaf_index), // element 3 (leaf index)
966                                    ]);
967                                    (key, value)
968                                })
969                                .collect();
970
971                            // Ensure unique keys to avoid duplicates
972                            let mut seen = alloc::collections::BTreeSet::new();
973                            let unique_entries: Vec<_> =
974                                entries.into_iter().filter(|(key, _)| seen.insert(*key)).collect();
975
976                            let tree = Smt::with_entries(unique_entries.clone()).unwrap();
977                            RandomSmt { tree, entries: unique_entries }
978                        },
979                    )
980                })
981                .boxed()
982        }
983    }
984
985    proptest! {
986        #![proptest_config(ProptestConfig::with_cases(20))]
987
988        #[test]
989        fn simple_smt_path_consistency(tree_data in any::<RandomSimpleSmt>()) {
990            let RandomSimpleSmt { tree, entries } = tree_data;
991
992            for (leaf_index, value) in &entries {
993                let merkle_path = tree.get_path(&LeafIndex::new(*leaf_index).unwrap());
994                let sparse_path = SparseMerklePath::from_sized_iter(merkle_path.clone().into_iter()).unwrap();
995
996                // Verify both paths have same depth
997                prop_assert_eq!(merkle_path.depth(), sparse_path.depth());
998
999                // Verify both paths produce same root for the same value
1000                let merkle_root = merkle_path.compute_root(*leaf_index, *value).unwrap();
1001                let sparse_root = sparse_path.compute_root(*leaf_index, *value).unwrap();
1002                prop_assert_eq!(merkle_root, sparse_root);
1003
1004                // Verify both paths verify correctly
1005                let tree_root = tree.root();
1006                prop_assert!(merkle_path.verify(*leaf_index, *value, &tree_root).is_ok());
1007                prop_assert!(sparse_path.verify(*leaf_index, *value, &tree_root).is_ok());
1008
1009                // Test with random additional leaf
1010                let random_leaf = Word::new([Felt::ONE; 4]);
1011                let random_index = *leaf_index ^ 1; // Ensure it's a sibling
1012
1013                // Both should fail verification with wrong leaf
1014                let merkle_wrong = merkle_path.verify(random_index, random_leaf, &tree_root);
1015                let sparse_wrong = sparse_path.verify(random_index, random_leaf, &tree_root);
1016                prop_assert_eq!(merkle_wrong.is_err(), sparse_wrong.is_err());
1017            }
1018        }
1019
1020        #[test]
1021        fn smt_path_consistency(tree_data in any::<RandomSmt>()) {
1022            let RandomSmt { tree, entries } = tree_data;
1023
1024            for (key, _value) in &entries {
1025                let (merkle_path, leaf) = tree.open(key).into_parts();
1026                let sparse_path = SparseMerklePath::from_sized_iter(merkle_path.clone().into_iter()).unwrap();
1027
1028                let leaf_index = Smt::key_to_leaf_index(key).value();
1029                let actual_value = leaf.hash(); // Use the actual leaf hash
1030
1031                // Verify both paths have same depth
1032                prop_assert_eq!(merkle_path.depth(), sparse_path.depth());
1033
1034                // Verify both paths produce same root for the same value
1035                let merkle_root = merkle_path.compute_root(leaf_index, actual_value).unwrap();
1036                let sparse_root = sparse_path.compute_root(leaf_index, actual_value).unwrap();
1037                prop_assert_eq!(merkle_root, sparse_root);
1038
1039                // Verify both paths verify correctly
1040                let tree_root = tree.root();
1041                prop_assert!(merkle_path.verify(leaf_index, actual_value, &tree_root).is_ok());
1042                prop_assert!(sparse_path.verify(leaf_index, actual_value, &tree_root).is_ok());
1043
1044                // Test authenticated nodes consistency
1045                let merkle_auth = merkle_path.authenticated_nodes(leaf_index, actual_value).unwrap().collect::<Vec<_>>();
1046                let sparse_auth = sparse_path.authenticated_nodes(leaf_index, actual_value).unwrap().collect::<Vec<_>>();
1047                prop_assert_eq!(merkle_auth, sparse_auth);
1048            }
1049        }
1050
1051        #[test]
1052        fn reverse_conversion_from_sparse(tree_data in any::<RandomMerkleTree>()) {
1053            let RandomMerkleTree { tree, leaves, indices } = tree_data;
1054
1055            for (i, &leaf_index) in indices.iter().enumerate() {
1056                let leaf = leaves[i];
1057                let merkle_path = tree.get_path(NodeIndex::new(tree.depth(), leaf_index).unwrap()).unwrap();
1058
1059                // Create SparseMerklePath first, then convert to MerklePath
1060                let sparse_path = SparseMerklePath::from_sized_iter(merkle_path.clone().into_iter()).unwrap();
1061                let converted_merkle = MerklePath::from(sparse_path.clone());
1062
1063                // Verify conversion back and forth works
1064                let back_to_sparse = SparseMerklePath::try_from(converted_merkle.clone()).unwrap();
1065                prop_assert_eq!(sparse_path, back_to_sparse);
1066
1067                // Verify all APIs work identically
1068                prop_assert_eq!(merkle_path.depth(), converted_merkle.depth());
1069
1070                let merkle_root = merkle_path.compute_root(leaf_index, leaf).unwrap();
1071                let converted_root = converted_merkle.compute_root(leaf_index, leaf).unwrap();
1072                prop_assert_eq!(merkle_root, converted_root);
1073            }
1074        }
1075    }
1076}