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    };
128}
129
130/// Macro for generating the range proof implementation
131#[macro_export]
132macro_rules! impl_range_proof_merkle_tree_scheme {
133    ($name: ident) => {
134        impl<E, H, I, const ARITY: usize, T> $crate::RangeProofMerkleTreeScheme
135            for $name<E, H, I, ARITY, T>
136        where
137            E: $crate::Element,
138            H: $crate::DigestAlgorithm<E, I, T>,
139            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
140            T: $crate::NodeValue,
141        {
142            type RangeMembershipProof = $crate::internal::MerkleTreeRangeProof<T>;
143
144            fn range_lookup(
145                &self,
146                start: impl Borrow<Self::Index>,
147                end: impl Borrow<Self::Index>,
148            ) -> LookupResult<
149                (
150                    ark_std::vec::Vec<Self::Index>,
151                    ark_std::vec::Vec<Self::Element>,
152                ),
153                Self::RangeMembershipProof,
154                (),
155            > {
156                let start = start.borrow();
157                let end = end.borrow();
158                let start_path = start.to_traversal_path(self.height);
159                let end_path = end.to_traversal_path(self.height);
160                self.root
161                    .range_lookup_internal(self.height, &start_path, &end_path, true, true)
162            }
163
164            fn verify_range_proof(
165                commitment: impl Borrow<Self::Commitment>,
166                indices: &[impl Borrow<Self::Index>],
167                elements: &[impl Borrow<Self::Element>],
168                proof: impl Borrow<Self::RangeMembershipProof>,
169            ) -> Result<VerificationResult, MerkleTreeError> {
170                $crate::internal::verify_merkle_range_proof::<E, H, I, ARITY, T>(
171                    commitment.borrow(),
172                    indices,
173                    elements,
174                    proof.borrow(),
175                )
176            }
177        }
178    };
179}
180
181/// Macro for generating a forgetable merkle tree implementation
182#[macro_export]
183macro_rules! impl_forgetable_merkle_tree_scheme {
184    ($name: ident) => {
185        impl<E, H, I, const ARITY: usize, T> $crate::ForgetableMerkleTreeScheme
186            for $name<E, H, I, ARITY, T>
187        where
188            E: $crate::Element,
189            H: $crate::DigestAlgorithm<E, I, T>,
190            I: $crate::Index + $crate::ToTraversalPath<ARITY>,
191            T: $crate::NodeValue,
192        {
193            fn from_commitment(
194                com: impl Borrow<Self::Commitment>,
195                height: usize,
196                num_leaves: u64,
197            ) -> Self {
198                let com = com.borrow();
199                $name {
200                    root: Arc::new($crate::internal::MerkleNode::ForgottenSubtree {
201                        value: com.clone(),
202                    }),
203                    height,
204                    num_leaves,
205                    _phantom: PhantomData,
206                }
207            }
208
209            fn forget(
210                &mut self,
211                pos: impl Borrow<Self::Index>,
212            ) -> LookupResult<Self::Element, Self::MembershipProof, ()> {
213                let pos = pos.borrow();
214                let traversal_path = pos.to_traversal_path(self.height);
215                let (new_root, result) = self.root.forget_internal(self.height, &traversal_path);
216                self.root = new_root;
217                match result {
218                    LookupResult::Ok(elem, proof) => LookupResult::Ok(elem, proof),
219                    LookupResult::NotInMemory => LookupResult::NotInMemory,
220                    LookupResult::NotFound(_) => LookupResult::NotFound(()),
221                }
222            }
223
224            fn remember(
225                &mut self,
226                pos: impl Borrow<Self::Index>,
227                element: impl Borrow<Self::Element>,
228                proof: impl Borrow<Self::MembershipProof>,
229            ) -> Result<(), MerkleTreeError> {
230                let pos = pos.borrow();
231                let element = element.borrow();
232                let proof = proof.borrow();
233                if Self::verify(&self.commitment(), pos, element, proof)?.is_err() {
234                    Err(MerkleTreeError::InconsistentStructureError(
235                        "Wrong proof".to_string(),
236                    ))
237                } else {
238                    let traversal_path = pos.to_traversal_path(self.height);
239                    self.root = self.root.remember_internal::<H, ARITY>(
240                        self.height,
241                        &traversal_path,
242                        pos,
243                        Some(element),
244                        proof.path_values(),
245                    )?;
246                    Ok(())
247                }
248            }
249        }
250    };
251}
252
253/// Macros for implementing ToTreversalPath for primitive types
254#[macro_export]
255macro_rules! impl_to_traversal_path_primitives {
256    ($t: ty) => {
257        impl<const ARITY: usize> $crate::ToTraversalPath<ARITY> for $t {
258            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
259                let mut pos = *self as u64;
260                let mut ret = vec![];
261                for _i in 0..height {
262                    ret.push((pos % (ARITY as u64)) as usize);
263                    pos /= ARITY as u64;
264                }
265                ret
266            }
267        }
268    };
269}
270
271/// Macros for implementing ToTreversalPath for BigUint types
272#[macro_export]
273macro_rules! impl_to_traversal_path_biguint {
274    ($t: ty) => {
275        impl<const ARITY: usize> $crate::ToTraversalPath<ARITY> for $t {
276            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
277                let mut pos: BigUint = <Self as Into<BigUint>>::into(self.clone());
278                let mut ret = vec![];
279                for _i in 0..height {
280                    ret.push((&pos % ARITY).to_usize().unwrap());
281                    pos /= ARITY;
282                }
283                ret
284            }
285        }
286    };
287}