Skip to main content

iris_ztd/
prove.rs

1use super::{Digest, Hashable, NounDecode, NounEncode};
2use alloc::vec::Vec;
3#[cfg(feature = "wasm")]
4use alloc::{boxed::Box, format, string::ToString};
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, NounEncode, NounDecode, Hashable, Serialize, Deserialize)]
8#[iris_ztd_derive::wasm_noun_codec]
9pub struct MerkleProof {
10    pub root: Digest,
11    pub path: Vec<Digest>,
12}
13
14#[derive(Debug, Clone, NounEncode, NounDecode, Hashable, Serialize, Deserialize)]
15#[iris_ztd_derive::wasm_noun_codec]
16pub struct MerkleProvenAxis {
17    pub proof: MerkleProof,
18    pub axis: u64,
19}
20
21impl MerkleProof {
22    /// Prove a 0-indexed leaf
23    ///
24    /// This is important - unlike `prove-hashable-by-index:merkle` the index here is 0-based, not 1.
25    pub fn prove_hashable<T: Hashable>(item: &T, index: usize) -> MerkleProvenAxis {
26        let Some((left, right)) = item.hashable_pair() else {
27            return MerkleProvenAxis {
28                proof: Self {
29                    root: item.hash(),
30                    path: Vec::new(),
31                },
32                axis: 1,
33            };
34        };
35        let lc = left.leaf_count();
36        if index < lc {
37            let mut rec = Self::prove_hashable(&left, index);
38            let sib = right.hash();
39            rec.proof.root = (rec.proof.root, sib).hash();
40            rec.proof.path.push(sib);
41            // This is like peg, but we invert the bit pattern, because we are implicitly flipping the leading 1 as well
42            let alz = rec.axis.leading_zeros();
43            rec.axis ^= 0b11 << (63 - alz);
44            rec
45        } else {
46            let mut rec = Self::prove_hashable(&right, index - lc);
47            let sib = left.hash();
48            rec.proof.root = (sib, rec.proof.root).hash();
49            rec.proof.path.push(sib);
50            let alz = rec.axis.leading_zeros();
51            rec.axis ^= 0b10 << (63 - alz);
52            rec
53        }
54    }
55
56    pub fn verify(&self, mut axis: u64, hashable: &impl Hashable) -> bool {
57        let mut leaf = hashable.hash();
58        let mut path = &self.path[..];
59
60        while axis > 1 {
61            let Some((sib, rest)) = path.split_first() else {
62                return false;
63            };
64            path = rest;
65            if axis.is_multiple_of(2) {
66                leaf = (leaf, sib).hash();
67            } else {
68                leaf = (sib, leaf).hash();
69            }
70            axis /= 2;
71        }
72
73        axis == 1 && self.root == leaf && path.is_empty()
74    }
75
76    pub fn visible_hashes(
77        &self,
78        mut axis: u64,
79        hashable: &impl Hashable,
80    ) -> Option<Vec<(u64, Digest)>> {
81        let mut hashes = Vec::new();
82        let mut leaf = hashable.hash();
83        let mut path = &self.path[..];
84
85        while axis > 1 {
86            let (sib, rest) = path.split_first()?;
87            path = rest;
88            if axis.is_multiple_of(2) {
89                hashes.push((axis ^ 1, *sib));
90                hashes.push((axis, leaf));
91                leaf = (leaf, sib).hash();
92            } else {
93                hashes.push((axis, leaf));
94                hashes.push((axis ^ 1, *sib));
95                leaf = (sib, leaf).hash();
96            }
97            axis /= 2;
98        }
99
100        if axis == 1 && self.root == leaf && path.is_empty() {
101            hashes.push((1, leaf));
102            Some(hashes)
103        } else {
104            None
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::HashableList;
113    use alloc::string::ToString;
114
115    #[test]
116    fn test_empty_proof() {
117        let MerkleProvenAxis { proof, axis } = MerkleProof::prove_hashable(&(), 0);
118        assert_eq!(axis, 1);
119        assert_eq!(proof.root.to_string(), ().hash().to_string());
120        assert_eq!(proof.path.len(), 0);
121        assert!(proof.verify(axis, &()));
122        let all_hashes = proof.visible_hashes(axis, &());
123        let all_hashes = all_hashes
124            .unwrap()
125            .iter()
126            .map(|(a, h)| (*a, h.to_string()))
127            .collect::<Vec<_>>();
128        assert_eq!(all_hashes, &[(1, ().hash().to_string())]);
129    }
130
131    #[test]
132    fn test_left_proof() {
133        let MerkleProvenAxis { proof, axis } = MerkleProof::prove_hashable(&((), ()), 0);
134        assert_eq!(axis, 2);
135        assert_eq!(
136            proof.root.to_string(),
137            "3LPSS51pUxLaxMD8VjyBSW6S9sotLpfx65zibBvm5k1xu18qt5ZGp3S"
138        );
139        assert_eq!(
140            proof.path.iter().map(|v| v.to_string()).collect::<Vec<_>>(),
141            ["3Ssr4tiWsbX5CE3AG6p5qPHP51fiyvtt1XEEHmSbGgDjp3qjUew6DFB"]
142        );
143        assert!(proof.verify(axis, &()));
144
145        let all_hashes = proof.visible_hashes(axis, &());
146        let all_hashes = all_hashes
147            .unwrap()
148            .iter()
149            .map(|(a, h)| (*a, h.to_string()))
150            .collect::<Vec<_>>();
151        assert_eq!(
152            all_hashes,
153            &[
154                (3, ().hash().to_string()),
155                (2, ().hash().to_string()),
156                (1, proof.root.to_string()),
157            ]
158        );
159    }
160
161    #[test]
162    fn test_right_proof() {
163        let MerkleProvenAxis { proof, axis } = MerkleProof::prove_hashable(&((), ()), 1);
164        assert_eq!(
165            proof.root.to_string(),
166            "3LPSS51pUxLaxMD8VjyBSW6S9sotLpfx65zibBvm5k1xu18qt5ZGp3S"
167        );
168        assert_eq!(
169            proof.path.iter().map(|v| v.to_string()).collect::<Vec<_>>(),
170            ["3Ssr4tiWsbX5CE3AG6p5qPHP51fiyvtt1XEEHmSbGgDjp3qjUew6DFB"]
171        );
172        assert_eq!(axis, 3);
173        assert!(proof.verify(axis, &()));
174
175        let all_hashes = proof.visible_hashes(axis, &());
176        let all_hashes = all_hashes
177            .unwrap()
178            .iter()
179            .map(|(a, h)| (*a, h.to_string()))
180            .collect::<Vec<_>>();
181        assert_eq!(
182            all_hashes,
183            &[
184                (3, ().hash().to_string()),
185                (2, ().hash().to_string()),
186                (1, proof.root.to_string()),
187            ]
188        );
189    }
190
191    #[test]
192    fn test_complex_proof() {
193        let MerkleProvenAxis { proof, axis } =
194            MerkleProof::prove_hashable(&((1u64, 2u64), (3u64, 4u64)), 2);
195        assert_eq!(
196            proof.root.to_string(),
197            "9BC9gRQaJ7Ub4SivF6NmPBQrmqfwdKeDSkbkRjmnKf9yYscct3AcohH"
198        );
199        assert_eq!(
200            proof.path.iter().map(|v| v.to_string()).collect::<Vec<_>>(),
201            [
202                "CdEJceqNNH5iCGYEsWhRf2gHE37zbJXVkVPLpfWW7uYrJjt8magUvgi",
203                "BqxDmSrtFP6QsDuoYxjaFxedEzGpy7gfwhmtZnD25FxeedB1ssNPH4t"
204            ]
205        );
206        assert_eq!(axis, 6);
207        assert!(proof.verify(axis, &3u64));
208
209        let all_hashes = proof.visible_hashes(axis, &3u64);
210        let all_hashes = all_hashes
211            .unwrap()
212            .iter()
213            .map(|(a, h)| (*a, h.to_string()))
214            .collect::<Vec<_>>();
215        assert_eq!(
216            all_hashes,
217            &[
218                (7, 4.hash().to_string()),
219                (6, 3.hash().to_string()),
220                (3, (3, 4).hash().to_string()),
221                (2, (1, 2).hash().to_string()),
222                (1, proof.root.to_string()),
223            ]
224        );
225    }
226
227    #[test]
228    fn test_list_proof() {
229        let lst = [0u64, 1u64];
230        let lst = HashableList(&lst[..]);
231        let MerkleProvenAxis { proof, axis } = MerkleProof::prove_hashable(&(&lst, ()), 0);
232        assert_eq!(axis, 2);
233        assert_eq!(
234            proof.root.to_string(),
235            "8cSTFCsmaL4KTMVq6RQSQaMMMQfb3YpT6xR1YmnRtG7P4WurnhRRDbM"
236        );
237        assert_eq!(
238            proof.path.iter().map(|v| v.to_string()).collect::<Vec<_>>(),
239            ["3Ssr4tiWsbX5CE3AG6p5qPHP51fiyvtt1XEEHmSbGgDjp3qjUew6DFB"]
240        );
241        assert!(proof.verify(axis, &lst));
242
243        let all_hashes = proof.visible_hashes(axis, &lst);
244        let all_hashes = all_hashes
245            .unwrap()
246            .iter()
247            .map(|(a, h)| (*a, h.to_string()))
248            .collect::<Vec<_>>();
249        assert_eq!(
250            all_hashes,
251            &[
252                (3, ().hash().to_string()),
253                (2, lst.hash().to_string()),
254                (1, proof.root.to_string()),
255            ]
256        );
257    }
258}