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