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