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