1use algebra::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 std::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 std::rc::Rc;
182
183    use crate::{
184        crh::{
185            pedersen::{constraints::PedersenCRHGadget, PedersenCRH, PedersenWindow},
186            FixedLengthCRH, FixedLengthCRHGadget,
187        },
188        merkle_tree::*,
189    };
190    use algebra::{curves::jubjub::JubJubAffine as JubJub, fields::jubjub::fq::Fq};
191    use r1cs_core::ConstraintSystem;
192    use rand::SeedableRng;
193    use rand_xorshift::XorShiftRng;
194
195    use super::*;
196    use r1cs_std::{
197        groups::curves::twisted_edwards::jubjub::JubJubGadget,
198        test_constraint_system::TestConstraintSystem,
199    };
200
201    #[derive(Clone)]
202    pub(super) struct Window4x256;
203    impl PedersenWindow for Window4x256 {
204        const WINDOW_SIZE: usize = 4;
205        const NUM_WINDOWS: usize = 256;
206    }
207
208    type H = PedersenCRH<JubJub, Window4x256>;
209    type HG = PedersenCRHGadget<JubJub, Fq, JubJubGadget>;
210
211    struct JubJubMerkleTreeParams;
212
213    impl MerkleTreeConfig for JubJubMerkleTreeParams {
214        const HEIGHT: usize = 32;
215        type H = H;
216    }
217
218    type JubJubMerkleTree = MerkleHashTree<JubJubMerkleTreeParams>;
219
220    fn generate_merkle_tree(leaves: &[[u8; 30]], use_bad_root: bool) -> () {
221        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
222
223        let crh_parameters = Rc::new(H::setup(&mut rng).unwrap());
224        let tree = JubJubMerkleTree::new(crh_parameters.clone(), leaves).unwrap();
225        let root = tree.root();
226        let mut satisfied = true;
227        for (i, leaf) in leaves.iter().enumerate() {
228            let mut cs = TestConstraintSystem::<Fq>::new();
229            let proof = tree.generate_proof(i, &leaf).unwrap();
230            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
231
232            // Allocate Merkle Tree Root
233            let root = <HG as FixedLengthCRHGadget<H, _>>::OutputGadget::alloc(
234                &mut cs.ns(|| format!("new_digest_{}", i)),
235                || {
236                    if use_bad_root {
237                        Ok(<H as FixedLengthCRH>::Output::default())
238                    } else {
239                        Ok(root)
240                    }
241                },
242            )
243            .unwrap();
244
245            let constraints_from_digest = cs.num_constraints();
246            println!("constraints from digest: {}", constraints_from_digest);
247
248            // Allocate Parameters for CRH
249            let crh_parameters = <HG as FixedLengthCRHGadget<H, Fq>>::ParametersGadget::alloc(
250                &mut cs.ns(|| format!("new_parameters_{}", i)),
251                || Ok(crh_parameters.clone()),
252            )
253            .unwrap();
254
255            let constraints_from_parameters = cs.num_constraints() - constraints_from_digest;
256            println!(
257                "constraints from parameters: {}",
258                constraints_from_parameters
259            );
260
261            // Allocate Leaf
262            let leaf_g = UInt8::constant_vec(leaf);
263
264            let constraints_from_leaf =
265                cs.num_constraints() - constraints_from_parameters - constraints_from_digest;
266            println!("constraints from leaf: {}", constraints_from_leaf);
267
268            // Allocate Merkle Tree Path
269            let cw = MerkleTreePathGadget::<_, HG, _>::alloc(
270                &mut cs.ns(|| format!("new_witness_{}", i)),
271                || Ok(proof),
272            )
273            .unwrap();
274
275            let constraints_from_path = cs.num_constraints()
276                - constraints_from_parameters
277                - constraints_from_digest
278                - constraints_from_leaf;
279            println!("constraints from path: {}", constraints_from_path);
280            let leaf_g: &[UInt8] = leaf_g.as_slice();
281            cw.check_membership(
282                &mut cs.ns(|| format!("new_witness_check_{}", i)),
283                &crh_parameters,
284                &root,
285                &leaf_g,
286            )
287            .unwrap();
288            if !cs.is_satisfied() {
289                satisfied = false;
290                println!(
291                    "Unsatisfied constraint: {}",
292                    cs.which_is_unsatisfied().unwrap()
293                );
294            }
295            let setup_constraints = constraints_from_leaf
296                + constraints_from_digest
297                + constraints_from_parameters
298                + constraints_from_path;
299            println!(
300                "number of constraints: {}",
301                cs.num_constraints() - setup_constraints
302            );
303        }
304
305        assert!(satisfied);
306    }
307
308    #[test]
309    fn good_root_test() {
310        let mut leaves = Vec::new();
311        for i in 0..4u8 {
312            let input = [i; 30];
313            leaves.push(input);
314        }
315        generate_merkle_tree(&leaves, false);
316    }
317
318    #[should_panic]
319    #[test]
320    fn bad_root_test() {
321        let mut leaves = Vec::new();
322        for i in 0..4u8 {
323            let input = [i; 30];
324            leaves.push(input);
325        }
326        generate_merkle_tree(&leaves, true);
327    }
328}