1use algebra_core::Field;
  2use r1cs_core::{Namespace, SynthesisError};
  3use r1cs_std::prelude::*;
  4
  5use crate::{
  6    crh::{FixedLengthCRH, FixedLengthCRHGadget},
  7    merkle_tree::*,
  8};
  9
 10use core::borrow::Borrow;
 11
 12pub struct PathVar<P, HGadget, ConstraintF>
 13where
 14    P: Config,
 15    HGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
 16    ConstraintF: Field,
 17{
 18    path: Vec<(HGadget::OutputVar, HGadget::OutputVar)>,
 19}
 20
 21impl<P, CRHGadget, ConstraintF> PathVar<P, CRHGadget, ConstraintF>
 22where
 23    P: Config,
 24    ConstraintF: Field,
 25    CRHGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
 26    <CRHGadget::OutputVar as R1CSVar<ConstraintF>>::Value: PartialEq,
 27{
 28    #[tracing::instrument(target = "r1cs", skip(self, parameters, root, leaf))]
 29    pub fn check_membership(
 30        &self,
 31        parameters: &CRHGadget::ParametersVar,
 32        root: &CRHGadget::OutputVar,
 33        leaf: impl ToBytesGadget<ConstraintF>,
 34    ) -> Result<Boolean<ConstraintF>, SynthesisError> {
 35        assert_eq!(self.path.len(), P::HEIGHT - 1);
 36        // Check that the hash of the given leaf matches the leaf hash in the membership
 37        // proof.
 38        let leaf_bits = leaf.to_bytes()?;
 39        let leaf_hash = CRHGadget::evaluate(parameters, &leaf_bits)?;
 40        let cs = leaf_hash.cs().or(root.cs()).unwrap();
 41
 42        // Check if leaf is one of the bottom-most siblings.
 43        let leaf_is_left = Boolean::new_witness(r1cs_core::ns!(cs, "leaf_is_left"), || {
 44            Ok(leaf_hash.value()?.eq(&self.path[0].0.value()?))
 45        })?;
 46
 47        let mut result =
 48            leaf_hash.is_eq(&leaf_is_left.select(&self.path[0].0, &self.path[0].1)?)?;
 49
 50        // Check levels between leaf level and root.
 51        let mut previous_hash = leaf_hash;
 52        for &(ref left_hash, ref right_hash) in &self.path {
 53            // Check if the previous_hash matches the correct current hash.
 54            let previous_is_left =
 55                Boolean::new_witness(r1cs_core::ns!(cs, "previous_is_left"), || {
 56                    Ok(previous_hash.value()?.eq(&left_hash.value()?))
 57                })?;
 58
 59            let ns = r1cs_core::ns!(cs, "enforcing that inner hash is correct");
 60            let equality_cmp = previous_is_left.select(left_hash, right_hash)?;
 61            result = result.and(&previous_hash.is_eq(&equality_cmp)?)?;
 62            drop(ns);
 63
 64            previous_hash =
 65                hash_inner_node::<P::H, CRHGadget, ConstraintF>(parameters, left_hash, right_hash)?;
 66        }
 67
 68        result.and(&root.is_eq(&previous_hash)?)
 69    }
 70}
 71
 72pub(crate) fn hash_inner_node<H, HG, ConstraintF>(
 73    parameters: &HG::ParametersVar,
 74    left_child: &HG::OutputVar,
 75    right_child: &HG::OutputVar,
 76) -> Result<HG::OutputVar, SynthesisError>
 77where
 78    ConstraintF: Field,
 79    H: FixedLengthCRH,
 80    HG: FixedLengthCRHGadget<H, ConstraintF>,
 81{
 82    let left_bytes = left_child.to_bytes()?;
 83    let right_bytes = right_child.to_bytes()?;
 84    let mut bytes = left_bytes;
 85    bytes.extend_from_slice(&right_bytes);
 86
 87    HG::evaluate(parameters, &bytes)
 88}
 89
 90impl<P, HGadget, ConstraintF> AllocVar<Path<P>, ConstraintF> for PathVar<P, HGadget, ConstraintF>
 91where
 92    P: Config,
 93    HGadget: FixedLengthCRHGadget<P::H, ConstraintF>,
 94    ConstraintF: Field,
 95{
 96    fn new_variable<T: Borrow<Path<P>>>(
 97        cs: impl Into<Namespace<ConstraintF>>,
 98        f: impl FnOnce() -> Result<T, SynthesisError>,
 99        mode: AllocationMode,
100    ) -> Result<Self, SynthesisError> {
101        let ns = cs.into();
102        let cs = ns.cs();
103        f().and_then(|val| {
104            let mut path = Vec::new();
105            for &(ref l, ref r) in val.borrow().path.iter() {
106                let l_hash = HGadget::OutputVar::new_variable(
107                    r1cs_core::ns!(cs, "l_child"),
108                    || Ok(l),
109                    mode,
110                )?;
111                let r_hash = HGadget::OutputVar::new_variable(
112                    r1cs_core::ns!(cs, "r_child"),
113                    || Ok(r),
114                    mode,
115                )?;
116                path.push((l_hash, r_hash));
117            }
118            Ok(PathVar { path })
119        })
120    }
121}
122
123#[cfg(test)]
124mod test {
125    use crate::{
126        crh::{
127            pedersen::{self, constraints::CRHGadget},
128            FixedLengthCRH, FixedLengthCRHGadget,
129        },
130        merkle_tree::*,
131    };
132    use algebra::ed_on_bls12_381::{EdwardsProjective as JubJub, Fq};
133    use r1cs_core::ConstraintSystem;
134    use rand::SeedableRng;
135    use rand_xorshift::XorShiftRng;
136
137    use super::*;
138    use r1cs_std::ed_on_bls12_381::EdwardsVar;
139
140    #[derive(Clone)]
141    pub(super) struct Window4x256;
142    impl pedersen::Window for Window4x256 {
143        const WINDOW_SIZE: usize = 4;
144        const NUM_WINDOWS: usize = 256;
145    }
146
147    type H = pedersen::CRH<JubJub, Window4x256>;
148    type HG = CRHGadget<JubJub, EdwardsVar, Window4x256>;
149
150    struct JubJubMerkleTreeParams;
151
152    impl Config for JubJubMerkleTreeParams {
153        const HEIGHT: usize = 32;
154        type H = H;
155    }
156
157    type JubJubMerkleTree = MerkleTree<JubJubMerkleTreeParams>;
158
159    fn generate_merkle_tree(leaves: &[[u8; 30]], use_bad_root: bool) -> () {
160        let mut rng = XorShiftRng::seed_from_u64(9174123u64);
161
162        let crh_parameters = H::setup(&mut rng).unwrap();
163        let tree = JubJubMerkleTree::new(crh_parameters.clone(), leaves).unwrap();
164        let root = tree.root();
165        let cs = ConstraintSystem::<Fq>::new_ref();
166        for (i, leaf) in leaves.iter().enumerate() {
167            let proof = tree.generate_proof(i, &leaf).unwrap();
168            assert!(proof.verify(&crh_parameters, &root, &leaf).unwrap());
169
170            // Allocate Merkle Tree Root
171            let root = <HG as FixedLengthCRHGadget<H, _>>::OutputVar::new_witness(
172                r1cs_core::ns!(cs, "new_digest"),
173                || {
174                    if use_bad_root {
175                        Ok(<H as FixedLengthCRH>::Output::default())
176                    } else {
177                        Ok(root)
178                    }
179                },
180            )
181            .unwrap();
182
183            let constraints_from_digest = cs.num_constraints();
184            println!("constraints from digest: {}", constraints_from_digest);
185
186            // Allocate Parameters for CRH
187            let crh_parameters = <HG as FixedLengthCRHGadget<H, Fq>>::ParametersVar::new_constant(
188                r1cs_core::ns!(cs, "new_parameter"),
189                &crh_parameters,
190            )
191            .unwrap();
192
193            let constraints_from_parameters = cs.num_constraints() - constraints_from_digest;
194            println!(
195                "constraints from parameters: {}",
196                constraints_from_parameters
197            );
198
199            // Allocate Leaf
200            let leaf_g = UInt8::constant_vec(leaf);
201
202            let constraints_from_leaf =
203                cs.num_constraints() - constraints_from_parameters - constraints_from_digest;
204            println!("constraints from leaf: {}", constraints_from_leaf);
205
206            // Allocate Merkle Tree Path
207            let cw =
208                PathVar::<_, HG, _>::new_witness(r1cs_core::ns!(cs, "new_witness"), || Ok(&proof))
209                    .unwrap();
210            for (i, (l, r)) in cw.path.iter().enumerate() {
211                assert_eq!(l.value().unwrap(), proof.path[i].0);
212                assert_eq!(r.value().unwrap(), proof.path[i].1);
213            }
214
215            let constraints_from_path = cs.num_constraints()
216                - constraints_from_parameters
217                - constraints_from_digest
218                - constraints_from_leaf;
219            println!("constraints from path: {}", constraints_from_path);
220            let leaf_g: &[_] = leaf_g.as_slice();
221            cw.check_membership(&crh_parameters, &root, &leaf_g)
222                .unwrap()
223                .enforce_equal(&Boolean::TRUE)
224                .unwrap();
225            let setup_constraints = constraints_from_leaf
226                + constraints_from_digest
227                + constraints_from_parameters
228                + constraints_from_path;
229            println!(
230                "number of constraints: {}",
231                cs.num_constraints() - setup_constraints
232            );
233        }
234
235        assert!(cs.is_satisfied().unwrap());
236    }
237
238    #[test]
239    fn good_root_test() {
240        let mut leaves = Vec::new();
241        for i in 0..4u8 {
242            let input = [i; 30];
243            leaves.push(input);
244        }
245        generate_merkle_tree(&leaves, false);
246    }
247
248    #[should_panic]
249    #[test]
250    fn bad_root_test() {
251        let mut leaves = Vec::new();
252        for i in 0..4u8 {
253            let input = [i; 30];
254            leaves.push(input);
255        }
256        generate_merkle_tree(&leaves, true);
257    }
258}