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