jf_merkle_tree/
macros.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//! Useful macros
8
9/// Macro for generating a standard merkle tree implementation
10#[macro_export]
11macro_rules! impl_merkle_tree_scheme {
12    ($name: ident) => {
13        /// A standard append only Merkle tree implementation
14        #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
15        #[serde(
16            bound = "E: ark_serialize::CanonicalSerialize + ark_serialize::CanonicalDeserialize,
17                     I: ark_serialize::CanonicalSerialize + ark_serialize::CanonicalDeserialize,"
18        )]
19        pub struct $name<E, H, I, const ARITY: usize, T>
20        where
21            E: $crate::Element,
22            H: $crate::DigestAlgorithm<E, I, T>,
23            I: $crate::Index,
24            T: $crate::NodeValue,
25        {
26            root: Arc<$crate::internal::MerkleNode<E, I, T>>,
27            height: usize,
28            num_leaves: u64,
29
30            _phantom: PhantomData<H>,
31        }
32
33        impl<E, H, I, const ARITY: usize, T> $crate::MerkleTreeScheme for $name<E, H, I, ARITY, T>
34        where
35            E: $crate::Element,
36            H: $crate::DigestAlgorithm<E, I, T>,
37            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
38            T: $crate::NodeValue,
39        {
40            type Element = E;
41            type Index = I;
42            type NodeValue = T;
43            type MembershipProof = $crate::internal::MerkleTreeProof<T>;
44            // TODO(Chengyu): implement batch membership proof
45            type BatchMembershipProof = ();
46            type Commitment = T;
47
48            const ARITY: usize = ARITY;
49
50            fn height(&self) -> usize {
51                self.height
52            }
53
54            fn capacity(&self) -> BigUint {
55                pow(BigUint::from(Self::ARITY), self.height)
56            }
57
58            fn num_leaves(&self) -> u64 {
59                self.num_leaves
60            }
61
62            fn commitment(&self) -> Self::Commitment {
63                self.root.value()
64            }
65
66            fn lookup(
67                &self,
68                pos: impl Borrow<Self::Index>,
69            ) -> LookupResult<&Self::Element, Self::MembershipProof, ()> {
70                let pos = pos.borrow();
71                let traversal_path = pos.to_traversal_path(self.height);
72                match self.root.lookup_internal(self.height, &traversal_path) {
73                    LookupResult::Ok(value, proof) => {
74                        LookupResult::Ok(&value, proof)
75                    },
76                    LookupResult::NotInMemory => LookupResult::NotInMemory,
77                    LookupResult::NotFound(_) => LookupResult::NotFound(()),
78                }
79            }
80
81            fn verify(
82                commitment: impl Borrow<Self::Commitment>,
83                pos: impl Borrow<Self::Index>,
84                element: impl Borrow<Self::Element>,
85                proof: impl Borrow<Self::MembershipProof>,
86            ) -> Result<VerificationResult, MerkleTreeError> {
87                $crate::internal::verify_merkle_proof::<E, H, I, ARITY, T>(commitment.borrow(), pos.borrow(), Some(element.borrow()), proof.borrow())
88            }
89
90            fn iter(&'_ self) -> $crate::MerkleTreeIter<'_, E, I, T> {
91                $crate::MerkleTreeIter::new(&self.root)
92            }
93        }
94
95        impl<'a, E, H, I, const ARITY: usize, T> IntoIterator for &'a $name<E, H, I, ARITY, T>
96        where
97            E: $crate::Element,
98            H: $crate::DigestAlgorithm<E, I, T>,
99            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
100            T: $crate::NodeValue,
101        {
102            type Item = (&'a I, &'a E);
103
104            type IntoIter = $crate::internal::MerkleTreeIter<'a, E, I, T>;
105
106            fn into_iter(self) -> Self::IntoIter {
107                $crate::internal::MerkleTreeIter::new(&self.root)
108            }
109        }
110
111        impl<E, H, I, const ARITY: usize, T> IntoIterator for $name<E, H, I, ARITY, T>
112        where
113            E: $crate::Element,
114            H: $crate::DigestAlgorithm<E, I, T>,
115            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
116            T: $crate::NodeValue,
117        {
118            type Item = (I, E);
119
120            type IntoIter = $crate::internal::MerkleTreeIntoIter<E, I, T>;
121
122            fn into_iter(self) -> Self::IntoIter {
123                $crate::internal::MerkleTreeIntoIter::new(self.root)
124            }
125        }
126
127        impl<E, H, I, const ARITY: usize, T> $name<E, H, I, ARITY, T>
128        where
129            E: $crate::Element,
130            H: $crate::DigestAlgorithm<E, I, T>,
131            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
132            T: $crate::NodeValue,
133        {
134            /// A helper function to collect all leaves with their corresponding proofs.
135            pub fn collect_leaves_with_proofs(&self) -> $crate::Vec<(&I, &E, $crate::internal::MerkleTreeProof<T>)> {
136                let mut collector = $crate::Vec::with_capacity(self.num_leaves() as usize);
137                let mut prealloc = $crate::vec![$crate::vec![T::default(); ARITY - 1]; self.height()];
138                self.root.collect_all_with_proofs(self.height(), &mut prealloc, &mut collector);
139                collector
140            }
141        }
142    };
143}
144
145/// Macro for generating the range proof implementation
146#[macro_export]
147macro_rules! impl_range_proof_merkle_tree_scheme {
148    ($name: ident) => {
149        impl<E, H, I, const ARITY: usize, T> $crate::RangeProofMerkleTreeScheme
150            for $name<E, H, I, ARITY, T>
151        where
152            E: $crate::Element,
153            H: $crate::DigestAlgorithm<E, I, T>,
154            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
155            T: $crate::NodeValue,
156        {
157            type RangeMembershipProof = $crate::internal::MerkleTreeRangeProof<T>;
158
159            fn range_lookup(
160                &self,
161                start: impl Borrow<Self::Index>,
162                end: impl Borrow<Self::Index>,
163            ) -> LookupResult<
164                (
165                    ark_std::vec::Vec<Self::Index>,
166                    ark_std::vec::Vec<Self::Element>,
167                ),
168                Self::RangeMembershipProof,
169                (),
170            > {
171                let start = start.borrow();
172                let end = end.borrow();
173                let start_path = start.to_traversal_path(self.height);
174                let end_path = end.to_traversal_path(self.height);
175                self.root
176                    .range_lookup_internal(self.height, &start_path, &end_path, true, true)
177            }
178
179            fn verify_range_proof(
180                commitment: impl Borrow<Self::Commitment>,
181                indices: &[impl Borrow<Self::Index>],
182                elements: &[impl Borrow<Self::Element>],
183                proof: impl Borrow<Self::RangeMembershipProof>,
184            ) -> Result<VerificationResult, MerkleTreeError> {
185                $crate::internal::verify_merkle_range_proof::<E, H, I, ARITY, T>(
186                    commitment.borrow(),
187                    indices,
188                    elements,
189                    proof.borrow(),
190                )
191            }
192        }
193    };
194}
195
196/// Macro for generating a forgetable merkle tree implementation
197#[macro_export]
198macro_rules! impl_forgetable_merkle_tree_scheme {
199    ($name: ident) => {
200        impl<E, H, I, const ARITY: usize, T> $crate::ForgetableMerkleTreeScheme
201            for $name<E, H, I, ARITY, T>
202        where
203            E: $crate::Element,
204            H: $crate::DigestAlgorithm<E, I, T>,
205            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
206            T: $crate::NodeValue,
207        {
208            fn from_commitment(
209                com: impl Borrow<Self::Commitment>,
210                height: usize,
211                num_leaves: u64,
212            ) -> Self {
213                let com = com.borrow();
214                $name {
215                    root: Arc::new($crate::internal::MerkleNode::ForgottenSubtree {
216                        value: com.clone(),
217                    }),
218                    height,
219                    num_leaves,
220                    _phantom: PhantomData,
221                }
222            }
223
224            fn forget(
225                &mut self,
226                pos: impl Borrow<Self::Index>,
227            ) -> LookupResult<Self::Element, Self::MembershipProof, ()> {
228                let pos = pos.borrow();
229                let traversal_path = pos.to_traversal_path(self.height);
230                let (new_root, result) = self.root.forget_internal(self.height, &traversal_path);
231                self.root = new_root;
232                match result {
233                    LookupResult::Ok(elem, proof) => LookupResult::Ok(elem, proof),
234                    LookupResult::NotInMemory => LookupResult::NotInMemory,
235                    LookupResult::NotFound(_) => LookupResult::NotFound(()),
236                }
237            }
238
239            fn remember(
240                &mut self,
241                pos: impl Borrow<Self::Index>,
242                element: impl Borrow<Self::Element>,
243                proof: impl Borrow<Self::MembershipProof>,
244            ) -> Result<(), MerkleTreeError> {
245                let pos = pos.borrow();
246                let element = element.borrow();
247                let proof = proof.borrow();
248                if Self::verify(&self.commitment(), pos, element, proof)?.is_err() {
249                    Err(MerkleTreeError::InconsistentStructureError(
250                        "Wrong proof".to_string(),
251                    ))
252                } else {
253                    let traversal_path = pos.to_traversal_path(self.height);
254                    self.root = self.root.remember_internal::<H, ARITY>(
255                        self.height,
256                        &traversal_path,
257                        pos,
258                        Some(element),
259                        proof.path_values(),
260                    )?;
261                    Ok(())
262                }
263            }
264        }
265    };
266}
267
268/// Macros for implementing ToTreversalPath for primitive types
269#[macro_export]
270macro_rules! impl_to_traversal_path_primitives {
271    ($t: ty) => {
272        impl<const ARITY: usize> $crate::ToTraversalPath<ARITY> for $t {
273            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
274                let mut pos = *self as u64;
275                let mut ret = vec![];
276                for _i in 0..height {
277                    ret.push((pos % (ARITY as u64)) as usize);
278                    pos /= ARITY as u64;
279                }
280                ret
281            }
282        }
283    };
284}
285
286/// Macros for implementing ToTreversalPath for BigUint types
287#[macro_export]
288macro_rules! impl_to_traversal_path_biguint {
289    ($t: ty) => {
290        impl<const ARITY: usize> $crate::ToTraversalPath<ARITY> for $t {
291            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
292                let mut pos: BigUint = <Self as Into<BigUint>>::into(self.clone());
293                let mut ret = vec![];
294                for _i in 0..height {
295                    ret.push((&pos % ARITY).to_usize().unwrap());
296                    pos /= ARITY;
297                }
298                ret
299            }
300        }
301    };
302}
jf_merkle_tree/macros.rs

jf_merkle_tree/
macros.rs
