1use algebra_core::Field;
  2use r1cs_core::{ConstraintSystem, SynthesisError};
  3use r1cs_std::{boolean::AllocatedBit, prelude::*};
  4
  5use crate::{
  6    crh::{FixedLengthCRH, FixedLengthCRHGadget},
  7    merkle_tree::*,
  8};
  9
 10use core::borrow::Borrow;
 11
 12pub struct MerkleTreePathGadget<P, HGadget, ConstraintF>
 13where
 14    P: MerkleTreeConfig,
 15    HGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
 16    ConstraintF: Field,
 17{
 18    path: Vec<(HGadget::OutputGadget, HGadget::OutputGadget)>,
 19}
 20
 21impl<P, CRHGadget, ConstraintF> MerkleTreePathGadget<P, CRHGadget, ConstraintF>
 22where
 23    P: MerkleTreeConfig,
 24    ConstraintF: Field,
 25    CRHGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
 26{
 27    pub fn check_membership<CS: ConstraintSystem<ConstraintF>>(
 28        &self,
 29        cs: CS,
 30        parameters: &CRHGadget::ParametersGadget,
 31        root: &CRHGadget::OutputGadget,
 32        leaf: impl ToBytesGadget<ConstraintF>,
 33    ) -> Result<(), SynthesisError> {
 34        self.conditionally_check_membership(cs, parameters, root, leaf, &Boolean::Constant(true))
 35    }
 36
 37    pub fn conditionally_check_membership<CS: ConstraintSystem<ConstraintF>>(
 38        &self,
 39        mut cs: CS,
 40        parameters: &CRHGadget::ParametersGadget,
 41        root: &CRHGadget::OutputGadget,
 42        leaf: impl ToBytesGadget<ConstraintF>,
 43        should_enforce: &Boolean,
 44    ) -> Result<(), SynthesisError> {
 45        assert_eq!(self.path.len(), P::HEIGHT - 1);
 46        // Check that the hash of the given leaf matches the leaf hash in the membership
 47        // proof.
 48        let leaf_bits = leaf.to_bytes(&mut cs.ns(|| "leaf_to_bytes"))?;
 49        let leaf_hash = CRHGadget::check_evaluation_gadget(
 50            cs.ns(|| "check_evaluation_gadget"),
 51            parameters,
 52            &leaf_bits,
 53        )?;
 54
 55        // Check if leaf is one of the bottom-most siblings.
 56        let leaf_is_left = AllocatedBit::alloc(&mut cs.ns(|| "leaf_is_left"), || {
 57            Ok(leaf_hash == self.path[0].0)
 58        })?
 59        .into();
 60        CRHGadget::OutputGadget::conditional_enforce_equal_or(
 61            &mut cs.ns(|| "check_leaf_is_left"),
 62            &leaf_is_left,
 63            &leaf_hash,
 64            &self.path[0].0,
 65            &self.path[0].1,
 66            should_enforce,
 67        )?;
 68
 69        // Check levels between leaf level and root.
 70        let mut previous_hash = leaf_hash;
 71        for (i, &(ref left_hash, ref right_hash)) in self.path.iter().enumerate() {
 72            // Check if the previous_hash matches the correct current hash.
 73            let previous_is_left =
 74                AllocatedBit::alloc(&mut cs.ns(|| format!("previous_is_left_{}", i)), || {
 75                    Ok(&previous_hash == left_hash)
 76                })?
 77                .into();
 78
 79            CRHGadget::OutputGadget::conditional_enforce_equal_or(
 80                &mut cs.ns(|| format!("check_equals_which_{}", i)),
 81                &previous_is_left,
 82                &previous_hash,
 83                left_hash,
 84                right_hash,
 85                should_enforce,
 86            )?;
 87
 88            previous_hash = hash_inner_node_gadget::<P::H, CRHGadget, ConstraintF, _>(
 89                &mut cs.ns(|| format!("hash_inner_node_{}", i)),
 90                parameters,
 91                left_hash,
 92                right_hash,
 93            )?;
 94        }
 95
 96        root.conditional_enforce_equal(
 97            &mut cs.ns(|| "root_is_last"),
 98            &previous_hash,
 99            should_enforce,
100        )
101    }
102}
103
104pub(crate) fn hash_inner_node_gadget<H, HG, ConstraintF, CS>(
105    mut cs: CS,
106    parameters: &HG::ParametersGadget,
107    left_child: &HG::OutputGadget,
108    right_child: &HG::OutputGadget,
109) -> Result<HG::OutputGadget, SynthesisError>
110where
111    ConstraintF: Field,
112    CS: ConstraintSystem<ConstraintF>,
113    H: FixedLengthCRH,
114    HG: FixedLengthCRHGadget<H, ConstraintF>,
115{
116    let left_bytes = left_child.to_bytes(&mut cs.ns(|| "left_to_bytes"))?;
117    let right_bytes = right_child.to_bytes(&mut cs.ns(|| "right_to_bytes"))?;
118    let mut bytes = left_bytes;
119    bytes.extend_from_slice(&right_bytes);
120
121    HG::check_evaluation_gadget(cs, parameters, &bytes)
122}
123
124impl<P, HGadget, ConstraintF> AllocGadget<MerkleTreePath<P>, ConstraintF>
125    for MerkleTreePathGadget<P, HGadget, ConstraintF>
126where
127    P: MerkleTreeConfig,
128    HGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
129    ConstraintF: Field,
130{
131    fn alloc<F, T, CS: ConstraintSystem<ConstraintF>>(
132        mut cs: CS,
133        value_gen: F,
134    ) -> Result<Self, SynthesisError>
135    where
136        F: FnOnce() -> Result<T, SynthesisError>,
137        T: Borrow<MerkleTreePath<P>>,
138    {
139        let mut path = Vec::new();
140        for (i, &(ref l, ref r)) in value_gen()?.borrow().path.iter().enumerate() {
141            let l_hash =
142                HGadget::OutputGadget::alloc(&mut cs.ns(|| format!("l_child_{}", i)), || {
143                    Ok(l.clone())
144                })?;
145            let r_hash =
146                HGadget::OutputGadget::alloc(&mut cs.ns(|| format!("r_child_{}", i)), || {
147                    Ok(r.clone())
148                })?;
149            path.push((l_hash, r_hash));
150        }
151        Ok(MerkleTreePathGadget { path })
152    }
153
154    fn alloc_input<F, T, CS: ConstraintSystem<ConstraintF>>(
155        mut cs: CS,
156        value_gen: F,
157    ) -> Result<Self, SynthesisError>
158    where
159        F: FnOnce() -> Result<T, SynthesisError>,
160        T: Borrow<MerkleTreePath<P>>,
161    {
162        let mut path = Vec::new();
163        for (i, &(ref l, ref r)) in value_gen()?.borrow().path.iter().enumerate() {
164            let l_hash = HGadget::OutputGadget::alloc_input(
165                &mut cs.ns(|| format!("l_child_{}", i)),
166                || Ok(l.clone()),
167            )?;
168            let r_hash = HGadget::OutputGadget::alloc_input(
169                &mut cs.ns(|| format!("r_child_{}", i)),
170                || Ok(r.clone()),
171            )?;
172            path.push((l_hash, r_hash));
173        }
174
175        Ok(MerkleTreePathGadget { path })
176    }
177}
178
179#[cfg(test)]
180mod test {
181    use crate::{
182        crh::{
183            pedersen::{constraints::PedersenCRHGadget, PedersenCRH, PedersenWindow},
184            FixedLengthCRH, FixedLengthCRHGadget,
185        },
186        merkle_tree::*,
187    };
188    use algebra::jubjub::{Fq, JubJubAffine as JubJub};
189    use r1cs_core::ConstraintSystem;
190    use rand::SeedableRng;
191    use rand_xorshift::XorShiftRng;
192
193    use super::*;
194    use r1cs_std::{jubjub::JubJubGadget, test_constraint_system::TestConstraintSystem};
195
196    #[derive(Clone)]
197    pub(super) struct Window4x256;
198    impl PedersenWindow for Window4x256 {
199        const WINDOW_SIZE: usize = 4;
200        const NUM_WINDOWS: usize = 256;
201    }
202
203    type H = PedersenCRH<JubJub, Window4x256>;
204    type HG = PedersenCRHGadget<JubJub, Fq, JubJubGadget>;
205
206    struct JubJubMerkleTreeParams;
207
208    impl MerkleTreeConfig for JubJubMerkleTreeParams {
209        const HEIGHT: usize = 32;
210        type H = H;
211    }
212
213    type JubJubMerkleTree = MerkleHashTree<JubJubMerkleTreeParams>;
214
215    fn generate_merkle_tree(leaves: &[[u8; 30]], use_bad_root: bool) -> () {
216        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
217
218        let crh_parameters = H::setup(&mut rng).unwrap();
219        let tree = JubJubMerkleTree::new(crh_parameters.clone(), leaves).unwrap();
220        let root = tree.root();
221        let mut satisfied = true;
222        for (i, leaf) in leaves.iter().enumerate() {
223            let mut cs = TestConstraintSystem::<Fq>::new();
224            let proof = tree.generate_proof(i, &leaf).unwrap();
225            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
226
227            // Allocate Merkle Tree Root
228            let root = <HG as FixedLengthCRHGadget<H, _>>::OutputGadget::alloc(
229                &mut cs.ns(|| format!("new_digest_{}", i)),
230                || {
231                    if use_bad_root {
232                        Ok(<H as FixedLengthCRH>::Output::default())
233                    } else {
234                        Ok(root)
235                    }
236                },
237            )
238            .unwrap();
239
240            let constraints_from_digest = cs.num_constraints();
241            println!("constraints from digest: {}", constraints_from_digest);
242
243            // Allocate Parameters for CRH
244            let crh_parameters = <HG as FixedLengthCRHGadget<H, Fq>>::ParametersGadget::alloc(
245                &mut cs.ns(|| format!("new_parameters_{}", i)),
246                || Ok(crh_parameters.clone()),
247            )
248            .unwrap();
249
250            let constraints_from_parameters = cs.num_constraints() - constraints_from_digest;
251            println!(
252                "constraints from parameters: {}",
253                constraints_from_parameters
254            );
255
256            // Allocate Leaf
257            let leaf_g = UInt8::constant_vec(leaf);
258
259            let constraints_from_leaf =
260                cs.num_constraints() - constraints_from_parameters - constraints_from_digest;
261            println!("constraints from leaf: {}", constraints_from_leaf);
262
263            // Allocate Merkle Tree Path
264            let cw = MerkleTreePathGadget::<_, HG, _>::alloc(
265                &mut cs.ns(|| format!("new_witness_{}", i)),
266                || Ok(proof),
267            )
268            .unwrap();
269
270            let constraints_from_path = cs.num_constraints()
271                - constraints_from_parameters
272                - constraints_from_digest
273                - constraints_from_leaf;
274            println!("constraints from path: {}", constraints_from_path);
275            let leaf_g: &[UInt8] = leaf_g.as_slice();
276            cw.check_membership(
277                &mut cs.ns(|| format!("new_witness_check_{}", i)),
278                &crh_parameters,
279                &root,
280                &leaf_g,
281            )
282            .unwrap();
283            if !cs.is_satisfied() {
284                satisfied = false;
285                println!(
286                    "Unsatisfied constraint: {}",
287                    cs.which_is_unsatisfied().unwrap()
288                );
289            }
290            let setup_constraints = constraints_from_leaf
291                + constraints_from_digest
292                + constraints_from_parameters
293                + constraints_from_path;
294            println!(
295                "number of constraints: {}",
296                cs.num_constraints() - setup_constraints
297            );
298        }
299
300        assert!(satisfied);
301    }
302
303    #[test]
304    fn good_root_test() {
305        let mut leaves = Vec::new();
306        for i in 0..4u8 {
307            let input = [i; 30];
308            leaves.push(input);
309        }
310        generate_merkle_tree(&leaves, false);
311    }
312
313    #[should_panic]
314    #[test]
315    fn bad_root_test() {
316        let mut leaves = Vec::new();
317        for i in 0..4u8 {
318            let input = [i; 30];
319            leaves.push(input);
320        }
321        generate_merkle_tree(&leaves, true);
322    }
323}