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: Element,
22            H: DigestAlgorithm<E, I, T>,
23            I: Index,
24            T: NodeValue,
25        {
26            root: Arc<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> MerkleTreeScheme for $name<E, H, I, ARITY, T>
34        where
35            E: Element,
36            H: DigestAlgorithm<E, I, T>,
37            I: Index + ToTraversalPath<ARITY>,
38            T: NodeValue,
39        {
40            type Element = E;
41            type Index = I;
42            type NodeValue = T;
43            type MembershipProof = 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().path_values())
88            }
89
90            fn iter(&self) -> MerkleTreeIter<E, I, T> {
91                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: Element,
98            H: DigestAlgorithm<E, I, T>,
99            I: Index + ToTraversalPath<ARITY>,
100            T: NodeValue,
101        {
102            type Item = (&'a I, &'a E);
103
104            type IntoIter = MerkleTreeIter<'a, E, I, T>;
105
106            fn into_iter(self) -> Self::IntoIter {
107                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: Element,
114            H: DigestAlgorithm<E, I, T>,
115            I: Index + ToTraversalPath<ARITY>,
116            T: NodeValue,
117        {
118            type Item = (I, E);
119
120            type IntoIter = MerkleTreeIntoIter<E, I, T>;
121
122            fn into_iter(self) -> Self::IntoIter {
123                MerkleTreeIntoIter::new(self.root)
124            }
125        }
126
127    };
128}
129
130/// Macro for generating a forgetable merkle tree implementation
131#[macro_export]
132macro_rules! impl_forgetable_merkle_tree_scheme {
133    ($name: ident) => {
134        impl<E, H, I, const ARITY: usize, T> ForgetableMerkleTreeScheme for $name<E, H, I, ARITY, T>
135        where
136            E: Element,
137            H: DigestAlgorithm<E, I, T>,
138            I: Index + ToTraversalPath<ARITY>,
139            T: NodeValue,
140        {
141            fn from_commitment(
142                com: impl Borrow<Self::Commitment>,
143                height: usize,
144                num_leaves: u64,
145            ) -> Self {
146                let com = com.borrow();
147                $name {
148                    root: Arc::new(MerkleNode::ForgottenSubtree { value: com.clone() }),
149                    height,
150                    num_leaves,
151                    _phantom: PhantomData,
152                }
153            }
154
155            fn forget(
156                &mut self,
157                pos: impl Borrow<Self::Index>,
158            ) -> LookupResult<Self::Element, Self::MembershipProof, ()> {
159                let pos = pos.borrow();
160                let traversal_path = pos.to_traversal_path(self.height);
161                let (new_root, result) = self.root.forget_internal(self.height, &traversal_path);
162                self.root = new_root;
163                match result {
164                    LookupResult::Ok(elem, proof) => LookupResult::Ok(elem, proof),
165                    LookupResult::NotInMemory => LookupResult::NotInMemory,
166                    LookupResult::NotFound(_) => LookupResult::NotFound(()),
167                }
168            }
169
170            fn remember(
171                &mut self,
172                pos: impl Borrow<Self::Index>,
173                element: impl Borrow<Self::Element>,
174                proof: impl Borrow<Self::MembershipProof>,
175            ) -> Result<(), MerkleTreeError> {
176                let pos = pos.borrow();
177                let element = element.borrow();
178                let proof = proof.borrow();
179                if Self::verify(&self.commitment(), pos, element, proof)?.is_err() {
180                    Err(MerkleTreeError::InconsistentStructureError(
181                        "Wrong proof".to_string(),
182                    ))
183                } else {
184                    let traversal_path = pos.to_traversal_path(self.height);
185                    self.root = self.root.remember_internal::<H, ARITY>(
186                        self.height,
187                        &traversal_path,
188                        pos,
189                        Some(element),
190                        proof.path_values(),
191                    )?;
192                    Ok(())
193                }
194            }
195        }
196    };
197}
198
199/// Macros for implementing ToTreversalPath for primitive types
200#[macro_export]
201macro_rules! impl_to_traversal_path_primitives {
202    ($t: ty) => {
203        impl<const ARITY: usize> ToTraversalPath<ARITY> for $t {
204            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
205                let mut pos = *self as u64;
206                let mut ret = vec![];
207                for _i in 0..height {
208                    ret.push((pos % (ARITY as u64)) as usize);
209                    pos /= ARITY as u64;
210                }
211                ret
212            }
213        }
214    };
215}
216
217/// Macros for implementing ToTreversalPath for BigUint types
218#[macro_export]
219macro_rules! impl_to_traversal_path_biguint {
220    ($t: ty) => {
221        impl<const ARITY: usize> ToTraversalPath<ARITY> for $t {
222            fn to_traversal_path(&self, height: usize) -> Vec<usize> {
223                let mut pos: BigUint = <Self as Into<BigUint>>::into(self.clone());
224                let mut ret = vec![];
225                for _i in 0..height {
226                    ret.push((&pos % ARITY).to_usize().unwrap());
227                    pos /= ARITY;
228                }
229                ret
230            }
231        }
232    };
233}