lean_imt/
stateless.rs

1//! Stateless proofs module.
2//!
3//! This is a feature module that provides functionalities to generate inclusion proofs without the tree data.
4
5use crate::lean_imt::LeanIMTError;
6
7/// Stateless proof element
8#[derive(Debug, Clone, PartialEq, Eq)]
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10pub struct ProofElement {
11    /// Tree level
12    level: usize,
13    /// Sibling index
14    sibling_index: usize,
15}
16
17impl ProofElement {
18    /// Creates a new `ProofElement` with the given level and sibling index.
19    pub fn new(level: usize, sibling_index: usize) -> Self {
20        Self {
21            level,
22            sibling_index,
23        }
24    }
25
26    /// Returns the level of the proof element.
27    pub fn level(&self) -> usize {
28        self.level
29    }
30
31    /// Returns the sibling index of the proof element.
32    pub fn sibling_index(&self) -> usize {
33        self.sibling_index
34    }
35
36    /// Determines if the original node was a right child.
37    /// If the sibling index is even, the original node was a right child.
38    pub fn is_right(&self) -> bool {
39        self.sibling_index & 1 == 0
40    }
41}
42
43/// Computes the set of elements of a Merkle proof of a Lean IMT without the tree data.
44///
45/// It receives the index of the leaf we want to prove and the tree size.
46///
47/// The set it returns will not contain the actual leaf data but their positions in the tree.
48pub fn stateless_path(
49    leaf_index: usize,
50    tree_size: usize,
51) -> Result<Vec<ProofElement>, LeanIMTError> {
52    if leaf_index >= tree_size {
53        return Err(LeanIMTError::IndexOutOfBounds);
54    }
55
56    let depth = tree_size.next_power_of_two().trailing_zeros() as usize;
57    let mut proof_elements = Vec::with_capacity(depth);
58    let (mut current_index, mut current_level_size) = (leaf_index, tree_size);
59
60    for level in 0..depth {
61        let is_right = (current_index & 1) != 0;
62        let sibling = if is_right {
63            current_index - 1
64        } else {
65            current_index + 1
66        };
67
68        if sibling < current_level_size {
69            proof_elements.push(ProofElement::new(level, sibling));
70        }
71
72        current_index /= 2;
73        current_level_size = (current_level_size + 1) / 2;
74    }
75
76    Ok(proof_elements)
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::{
83        hashed_tree::{HashedLeanIMT, LeanIMTHasher},
84        lean_imt::MerkleProof,
85    };
86    use rand::Rng;
87    use std::hash::{DefaultHasher, Hash, Hasher};
88
89    struct SampleHasher;
90
91    impl LeanIMTHasher<32> for SampleHasher {
92        fn hash(input: &[u8]) -> [u8; 32] {
93            let mut hasher = DefaultHasher::new();
94            input.hash(&mut hasher);
95            let h = hasher.finish();
96            let mut result = [0u8; 32];
97            result[..8].copy_from_slice(&h.to_le_bytes());
98            result
99        }
100    }
101
102    /// Reconstructs a Merkle proof from tree data and a slice of proof elements.
103    fn proof_from_indices(
104        tree: &HashedLeanIMT<32, SampleHasher>,
105        leaf_index: usize,
106        proof_elements: &[ProofElement],
107    ) -> MerkleProof<32> {
108        let mut siblings = Vec::with_capacity(proof_elements.len());
109        let mut bits = Vec::with_capacity(proof_elements.len());
110
111        for elem in proof_elements.iter() {
112            siblings.push(tree.get_node(elem.level(), elem.sibling_index()).unwrap());
113            bits.push(if elem.is_right() { 1 } else { 0 });
114        }
115
116        let mut encoded_index = 0;
117        for bit in bits.into_iter().rev() {
118            encoded_index = (encoded_index << 1) | bit;
119        }
120
121        MerkleProof {
122            root: tree.root().unwrap(),
123            leaf: tree.get_leaf(leaf_index).unwrap(),
124            siblings,
125            index: encoded_index,
126        }
127    }
128
129    // Helper function to create a random leaf
130    fn random_leaf(rng: &mut impl Rng) -> [u8; 32] {
131        let mut leaf = [0u8; 32];
132        rng.fill(&mut leaf);
133        leaf
134    }
135
136    #[test]
137    fn test_stateless_merkle_path() {
138        let mut rng = rand::rng();
139        let max_tree_size = 2_usize.pow(10);
140
141        let mut random_leaf_set = Vec::with_capacity(max_tree_size);
142        for _ in 0..max_tree_size {
143            random_leaf_set.push(random_leaf(&mut rng));
144        }
145
146        for size in 0..max_tree_size {
147            let tree =
148                HashedLeanIMT::<32, SampleHasher>::new(&random_leaf_set[0..size], SampleHasher)
149                    .unwrap();
150
151            for leaf_index in 0..size {
152                let stateless_indices = stateless_path(leaf_index, size).unwrap();
153                let stateless_proof = proof_from_indices(&tree, leaf_index, &stateless_indices);
154                let actual_proof = tree.generate_proof(leaf_index).unwrap();
155                assert_eq!(actual_proof, stateless_proof);
156            }
157        }
158    }
159}
lean_imt/stateless.rs

lean_imt/
stateless.rs
