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