1use algebra_core::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_constant<T, CS: ConstraintSystem<ConstraintF>>(
132        mut cs: CS,
133        val: T,
134    ) -> Result<Self, SynthesisError>
135    where
136        T: Borrow<MerkleTreePath<P>>,
137    {
138        let mut path = Vec::new();
139        for (i, &(ref l, ref r)) in val.borrow().path.iter().enumerate() {
140            let l_hash = HGadget::OutputGadget::alloc_constant(
141                &mut cs.ns(|| format!("l_child_{}", i)),
142                l.clone(),
143            )?;
144            let r_hash = HGadget::OutputGadget::alloc_constant(
145                &mut cs.ns(|| format!("r_child_{}", i)),
146                r.clone(),
147            )?;
148            path.push((l_hash, r_hash));
149        }
150        Ok(MerkleTreePathGadget { path })
151    }
152
153    fn alloc<F, T, CS: ConstraintSystem<ConstraintF>>(
154        mut cs: CS,
155        value_gen: F,
156    ) -> Result<Self, SynthesisError>
157    where
158        F: FnOnce() -> Result<T, SynthesisError>,
159        T: Borrow<MerkleTreePath<P>>,
160    {
161        let mut path = Vec::new();
162        for (i, &(ref l, ref r)) in value_gen()?.borrow().path.iter().enumerate() {
163            let l_hash =
164                HGadget::OutputGadget::alloc(&mut cs.ns(|| format!("l_child_{}", i)), || {
165                    Ok(l.clone())
166                })?;
167            let r_hash =
168                HGadget::OutputGadget::alloc(&mut cs.ns(|| format!("r_child_{}", i)), || {
169                    Ok(r.clone())
170                })?;
171            path.push((l_hash, r_hash));
172        }
173        Ok(MerkleTreePathGadget { path })
174    }
175
176    fn alloc_input<F, T, CS: ConstraintSystem<ConstraintF>>(
177        mut cs: CS,
178        value_gen: F,
179    ) -> Result<Self, SynthesisError>
180    where
181        F: FnOnce() -> Result<T, SynthesisError>,
182        T: Borrow<MerkleTreePath<P>>,
183    {
184        let mut path = Vec::new();
185        for (i, &(ref l, ref r)) in value_gen()?.borrow().path.iter().enumerate() {
186            let l_hash = HGadget::OutputGadget::alloc_input(
187                &mut cs.ns(|| format!("l_child_{}", i)),
188                || Ok(l.clone()),
189            )?;
190            let r_hash = HGadget::OutputGadget::alloc_input(
191                &mut cs.ns(|| format!("r_child_{}", i)),
192                || Ok(r.clone()),
193            )?;
194            path.push((l_hash, r_hash));
195        }
196
197        Ok(MerkleTreePathGadget { path })
198    }
199}
200
201#[cfg(test)]
202mod test {
203    use crate::{
204        crh::{
205            pedersen::{constraints::PedersenCRHGadget, PedersenCRH, PedersenWindow},
206            FixedLengthCRH, FixedLengthCRHGadget,
207        },
208        merkle_tree::*,
209    };
210    use algebra::edwards_on_bls12_381::{EdwardsAffine as JubJub, Fq};
211    use r1cs_core::ConstraintSystem;
212    use rand::SeedableRng;
213    use rand_xorshift::XorShiftRng;
214
215    use super::*;
216    use r1cs_std::{
217        edwards_on_bls12_381::EdwardsGadget, test_constraint_system::TestConstraintSystem,
218    };
219
220    #[derive(Clone)]
221    pub(super) struct Window4x256;
222    impl PedersenWindow for Window4x256 {
223        const WINDOW_SIZE: usize = 4;
224        const NUM_WINDOWS: usize = 256;
225    }
226
227    type H = PedersenCRH<JubJub, Window4x256>;
228    type HG = PedersenCRHGadget<JubJub, Fq, EdwardsGadget>;
229
230    struct JubJubMerkleTreeParams;
231
232    impl MerkleTreeConfig for JubJubMerkleTreeParams {
233        const HEIGHT: usize = 32;
234        type H = H;
235    }
236
237    type JubJubMerkleTree = MerkleHashTree<JubJubMerkleTreeParams>;
238
239    fn generate_merkle_tree(leaves: &[[u8; 30]], use_bad_root: bool) -> () {
240        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
241
242        let crh_parameters = H::setup(&mut rng).unwrap();
243        let tree = JubJubMerkleTree::new(crh_parameters.clone(), leaves).unwrap();
244        let root = tree.root();
245        let mut satisfied = true;
246        for (i, leaf) in leaves.iter().enumerate() {
247            let mut cs = TestConstraintSystem::<Fq>::new();
248            let proof = tree.generate_proof(i, &leaf).unwrap();
249            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
250
251            // Allocate Merkle Tree Root
252            let root = <HG as FixedLengthCRHGadget<H, _>>::OutputGadget::alloc(
253                &mut cs.ns(|| format!("new_digest_{}", i)),
254                || {
255                    if use_bad_root {
256                        Ok(<H as FixedLengthCRH>::Output::default())
257                    } else {
258                        Ok(root)
259                    }
260                },
261            )
262            .unwrap();
263
264            let constraints_from_digest = cs.num_constraints();
265            println!("constraints from digest: {}", constraints_from_digest);
266
267            // Allocate Parameters for CRH
268            let crh_parameters = <HG as FixedLengthCRHGadget<H, Fq>>::ParametersGadget::alloc(
269                &mut cs.ns(|| format!("new_parameters_{}", i)),
270                || Ok(crh_parameters.clone()),
271            )
272            .unwrap();
273
274            let constraints_from_parameters = cs.num_constraints() - constraints_from_digest;
275            println!(
276                "constraints from parameters: {}",
277                constraints_from_parameters
278            );
279
280            // Allocate Leaf
281            let leaf_g = UInt8::constant_vec(leaf);
282
283            let constraints_from_leaf =
284                cs.num_constraints() - constraints_from_parameters - constraints_from_digest;
285            println!("constraints from leaf: {}", constraints_from_leaf);
286
287            // Allocate Merkle Tree Path
288            let cw = MerkleTreePathGadget::<_, HG, _>::alloc(
289                &mut cs.ns(|| format!("new_witness_{}", i)),
290                || Ok(proof),
291            )
292            .unwrap();
293
294            let constraints_from_path = cs.num_constraints()
295                - constraints_from_parameters
296                - constraints_from_digest
297                - constraints_from_leaf;
298            println!("constraints from path: {}", constraints_from_path);
299            let leaf_g: &[UInt8] = leaf_g.as_slice();
300            cw.check_membership(
301                &mut cs.ns(|| format!("new_witness_check_{}", i)),
302                &crh_parameters,
303                &root,
304                &leaf_g,
305            )
306            .unwrap();
307            if !cs.is_satisfied() {
308                satisfied = false;
309                println!(
310                    "Unsatisfied constraint: {}",
311                    cs.which_is_unsatisfied().unwrap()
312                );
313            }
314            let setup_constraints = constraints_from_leaf
315                + constraints_from_digest
316                + constraints_from_parameters
317                + constraints_from_path;
318            println!(
319                "number of constraints: {}",
320                cs.num_constraints() - setup_constraints
321            );
322        }
323
324        assert!(satisfied);
325    }
326
327    #[test]
328    fn good_root_test() {
329        let mut leaves = Vec::new();
330        for i in 0..4u8 {
331            let input = [i; 30];
332            leaves.push(input);
333        }
334        generate_merkle_tree(&leaves, false);
335    }
336
337    #[should_panic]
338    #[test]
339    fn bad_root_test() {
340        let mut leaves = Vec::new();
341        for i in 0..4u8 {
342            let input = [i; 30];
343            leaves.push(input);
344        }
345        generate_merkle_tree(&leaves, true);
346    }
347}