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