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