jf_merkle_tree/
light_weight.rs

1// Copyright (c) 2022 Espresso Systems (espressosys.com)
2// This file is part of the Jellyfish library.
3
4// You should have received a copy of the MIT License
5// along with the Jellyfish library. If not, see <https://mit-license.org/>.
6
7//! A light weight merkle tree is an append only merkle tree who only keeps its
8//! frontier -- the right-most path.
9
10use super::{
11    internal::{
12        build_light_weight_tree_internal, MerkleNode, MerkleTreeIntoIter, MerkleTreeIter,
13        MerkleTreeProof,
14    },
15    AppendableMerkleTreeScheme, DigestAlgorithm, Element, ForgetableMerkleTreeScheme, Index,
16    LookupResult, MerkleProof, MerkleTreeScheme, NodeValue, ToTraversalPath,
17};
18use crate::{
19    errors::MerkleTreeError, impl_forgetable_merkle_tree_scheme, impl_merkle_tree_scheme,
20    VerificationResult,
21};
22use alloc::sync::Arc;
23use ark_std::{borrow::Borrow, fmt::Debug, marker::PhantomData, string::ToString, vec, vec::Vec};
24use num_bigint::BigUint;
25use num_traits::pow::pow;
26use serde::{Deserialize, Serialize};
27
28impl_merkle_tree_scheme!(LightWeightMerkleTree);
29impl_forgetable_merkle_tree_scheme!(LightWeightMerkleTree);
30
31impl<E, H, I, const ARITY: usize, T> LightWeightMerkleTree<E, H, I, ARITY, T>
32where
33    E: Element,
34    H: DigestAlgorithm<E, I, T>,
35    I: Index,
36    T: NodeValue,
37{
38    /// Initialize an empty Merkle tree.
39    pub fn new(height: usize) -> Self {
40        Self {
41            root: Arc::new(MerkleNode::<E, I, T>::Empty),
42            height,
43            num_leaves: 0,
44            _phantom: PhantomData,
45        }
46    }
47}
48
49impl<E, H, const ARITY: usize, T> LightWeightMerkleTree<E, H, u64, ARITY, T>
50where
51    E: Element,
52    H: DigestAlgorithm<E, u64, T>,
53    T: NodeValue,
54{
55    /// Construct a new Merkle tree with given height from a data slice
56    /// * `height` - height of the Merkle tree, if `None`, it will calculate the
57    ///   minimum height that could hold all elements.
58    /// * `elems` - an iterator to all elements
59    /// * `returns` - A constructed Merkle tree, or `Err()` if errors
60    pub fn from_elems(
61        height: Option<usize>,
62        elems: impl IntoIterator<Item = impl Borrow<E>>,
63    ) -> Result<Self, MerkleTreeError> {
64        let (root, height, num_leaves) =
65            build_light_weight_tree_internal::<E, H, ARITY, T>(height, elems)?;
66        Ok(Self {
67            root,
68            height,
69            num_leaves,
70            _phantom: PhantomData,
71        })
72    }
73}
74
75impl<E, H, const ARITY: usize, T> AppendableMerkleTreeScheme
76    for LightWeightMerkleTree<E, H, u64, ARITY, T>
77where
78    E: Element,
79    H: DigestAlgorithm<E, u64, T>,
80    T: NodeValue,
81{
82    fn push(&mut self, elem: impl Borrow<Self::Element>) -> Result<(), MerkleTreeError> {
83        <Self as AppendableMerkleTreeScheme>::extend(self, [elem])
84    }
85
86    fn extend(
87        &mut self,
88        elems: impl IntoIterator<Item = impl Borrow<Self::Element>>,
89    ) -> Result<(), MerkleTreeError> {
90        let mut iter = elems.into_iter().peekable();
91
92        let traversal_path =
93            ToTraversalPath::<ARITY>::to_traversal_path(&self.num_leaves, self.height);
94        let (root, num_inserted) = self.root.extend_and_forget_internal::<H, ARITY>(
95            self.height,
96            &self.num_leaves,
97            &traversal_path,
98            true,
99            &mut iter,
100        )?;
101        self.root = root;
102        self.num_leaves += num_inserted;
103        if iter.peek().is_some() {
104            return Err(MerkleTreeError::ExceedCapacity);
105        }
106        Ok(())
107    }
108}
109
110#[cfg(test)]
111mod mt_tests {
112    use crate::{
113        internal::MerkleNode,
114        prelude::{RescueLightWeightMerkleTree, RescueMerkleTree},
115        *,
116    };
117    use ark_bls12_377::Fr as Fr377;
118    use ark_bls12_381::Fr as Fr381;
119    use ark_bn254::Fr as Fr254;
120    use jf_rescue::RescueParameter;
121
122    #[test]
123    fn test_light_mt_builder() {
124        test_light_mt_builder_helper::<Fr254>();
125        test_light_mt_builder_helper::<Fr377>();
126        test_light_mt_builder_helper::<Fr381>();
127    }
128
129    fn test_light_mt_builder_helper<F: RescueParameter>() {
130        let arity: usize = RescueLightWeightMerkleTree::<F>::ARITY;
131        let mut data = vec![F::from(0u64); arity];
132        assert!(RescueLightWeightMerkleTree::<F>::from_elems(Some(1), &data).is_ok());
133        data.push(F::from(0u64));
134        assert!(RescueLightWeightMerkleTree::<F>::from_elems(Some(1), &data).is_err());
135    }
136
137    #[test]
138    fn test_light_mt_insertion() {
139        test_light_mt_insertion_helper::<Fr254>();
140        test_light_mt_insertion_helper::<Fr377>();
141        test_light_mt_insertion_helper::<Fr381>();
142    }
143
144    fn test_light_mt_insertion_helper<F: RescueParameter>() {
145        let mut mt = RescueLightWeightMerkleTree::<F>::new(2);
146        assert_eq!(mt.capacity(), BigUint::from(9u64));
147        assert!(mt.push(F::from(2u64)).is_ok());
148        assert!(mt.push(F::from(3u64)).is_ok());
149        assert!(mt.extend(&[F::from(0u64); 9]).is_err()); // Will err, but first 7 items will be inserted
150        assert_eq!(mt.num_leaves(), 9); // full merkle tree
151
152        // Now unable to insert more data
153        assert!(mt.push(F::from(0u64)).is_err());
154        assert!(mt.extend(&[]).is_ok());
155        assert!(mt.extend(&[F::from(1u64)]).is_err());
156
157        // Checks that the prior elements are all forgotten
158        (0..8).for_each(|i| assert!(mt.lookup(i).expect_not_in_memory().is_ok()));
159        assert!(mt.lookup(8).expect_ok().is_ok());
160    }
161
162    #[test]
163    fn test_light_mt_lookup() {
164        test_light_mt_lookup_helper::<Fr254>();
165        test_light_mt_lookup_helper::<Fr377>();
166        test_light_mt_lookup_helper::<Fr381>();
167    }
168
169    fn test_light_mt_lookup_helper<F: RescueParameter>() {
170        // singleton merkle tree test (#499)
171        let mt = RescueLightWeightMerkleTree::<F>::from_elems(None, [F::from(0u64)]).unwrap();
172        let (elem, _) = mt.lookup(0).expect_ok().unwrap();
173        assert_eq!(elem, &F::from(0u64));
174
175        let mut mt =
176            RescueLightWeightMerkleTree::<F>::from_elems(Some(2), [F::from(3u64), F::from(1u64)])
177                .unwrap();
178        let mut full_mt =
179            RescueMerkleTree::<F>::from_elems(Some(2), [F::from(3u64), F::from(1u64)]).unwrap();
180        assert!(mt.lookup(0).expect_not_in_memory().is_ok());
181        assert!(mt.lookup(1).expect_ok().is_ok());
182        assert!(mt.extend(&[F::from(33u64), F::from(41u64)]).is_ok());
183        assert!(full_mt.extend(&[F::from(33u64), F::from(41u64)]).is_ok());
184        assert!(mt.lookup(0).expect_not_in_memory().is_ok());
185        assert!(mt.lookup(1).expect_not_in_memory().is_ok());
186        assert!(mt.lookup(2).expect_not_in_memory().is_ok());
187        assert!(mt.lookup(3).expect_ok().is_ok());
188
189        // Should have the same commitment
190        assert_eq!(mt.commitment(), full_mt.commitment());
191
192        let commitment = mt.commitment();
193        let (elem, proof) = full_mt.lookup(0).expect_ok().unwrap();
194        assert_eq!(elem, &F::from(3u64));
195        assert!(
196            RescueLightWeightMerkleTree::<F>::verify(&commitment, 0, elem, &proof)
197                .unwrap()
198                .is_ok()
199        );
200
201        // Wrong element value, should fail.
202        assert!(
203            RescueLightWeightMerkleTree::<F>::verify(&commitment, 0, F::from(14u64), &proof)
204                .unwrap()
205                .is_err()
206        );
207
208        // Wrong pos, should fail.
209        assert!(
210            RescueLightWeightMerkleTree::<F>::verify(&commitment, 2, elem, &proof)
211                .unwrap()
212                .is_err()
213        );
214
215        let mut bad_proof = proof.clone();
216        bad_proof.0[0][0] = F::one();
217
218        assert!(
219            RescueLightWeightMerkleTree::<F>::verify(&commitment, 0, elem, &bad_proof)
220                .unwrap()
221                .is_err()
222        );
223    }
224
225    #[test]
226    fn test_light_mt_serde() {
227        test_light_mt_serde_helper::<Fr254>();
228        test_light_mt_serde_helper::<Fr377>();
229        test_light_mt_serde_helper::<Fr381>();
230    }
231
232    fn test_light_mt_serde_helper<F: RescueParameter>() {
233        let mt =
234            RescueLightWeightMerkleTree::<F>::from_elems(Some(2), [F::from(3u64), F::from(1u64)])
235                .unwrap();
236        let (_, proof) = mt.lookup(1).expect_ok().unwrap();
237
238        assert_eq!(
239            mt,
240            bincode::deserialize(&bincode::serialize(&mt).unwrap()).unwrap()
241        );
242        assert_eq!(
243            proof,
244            bincode::deserialize(&bincode::serialize(&proof).unwrap()).unwrap()
245        );
246    }
247}
jf_merkle_tree/light_weight.rs

jf_merkle_tree/
light_weight.rs
