sapling_crypto_ce/circuit/test/
mod.rs

1use bellman::pairing::{
2    Engine,
3};
4
5use bellman::pairing::ff::{
6    Field,
7    PrimeField,
8    PrimeFieldRepr
9};
10
11use bellman::{
12    LinearCombination,
13    SynthesisError,
14    ConstraintSystem,
15    Variable,
16    Index
17};
18
19use std::collections::HashMap;
20use std::fmt::Write;
21
22use byteorder::{BigEndian, ByteOrder};
23use std::cmp::Ordering;
24use std::collections::BTreeMap;
25use std::collections::HashSet;
26
27use blake2_rfc::blake2s::Blake2s;
28
29#[derive(Debug)]
30enum NamedObject {
31    Constraint(usize),
32    Var(Variable),
33    Namespace
34}
35
36/// Constraint system for testing purposes.
37pub struct TestConstraintSystem<E: Engine> {
38    named_objects: HashMap<String, NamedObject>,
39    current_namespace: Vec<String>,
40    constraints: Vec<(
41        LinearCombination<E>,
42        LinearCombination<E>,
43        LinearCombination<E>,
44        String
45    )>,
46    inputs: Vec<(E::Fr, String)>,
47    aux: Vec<(E::Fr, String)>
48}
49
50#[derive(Clone, Copy)]
51struct OrderedVariable(Variable);
52
53impl Eq for OrderedVariable {}
54impl PartialEq for OrderedVariable {
55    fn eq(&self, other: &OrderedVariable) -> bool {
56        match (self.0.get_unchecked(), other.0.get_unchecked()) {
57            (Index::Input(ref a), Index::Input(ref b)) => a == b,
58            (Index::Aux(ref a), Index::Aux(ref b)) => a == b,
59            _ => false
60        }
61    }
62}
63impl PartialOrd for OrderedVariable {
64    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
65        Some(self.cmp(other))
66    }
67}
68impl Ord for OrderedVariable {
69    fn cmp(&self, other: &Self) -> Ordering {
70        match (self.0.get_unchecked(), other.0.get_unchecked()) {
71            (Index::Input(ref a), Index::Input(ref b)) => a.cmp(b),
72            (Index::Aux(ref a), Index::Aux(ref b)) => a.cmp(b),
73            (Index::Input(_), Index::Aux(_)) => Ordering::Less,
74            (Index::Aux(_), Index::Input(_)) => Ordering::Greater
75        }
76    }
77}
78
79fn proc_lc<E: Engine>(
80    terms: &[(Variable, E::Fr)],
81) -> BTreeMap<OrderedVariable, E::Fr>
82{
83    let mut map = BTreeMap::new();
84    for &(var, coeff) in terms {
85        map.entry(OrderedVariable(var))
86           .or_insert(E::Fr::zero())
87           .add_assign(&coeff);
88    }
89
90    // Remove terms that have a zero coefficient to normalize
91    let mut to_remove = vec![];
92    for (var, coeff) in map.iter() {
93        if coeff.is_zero() {
94            to_remove.push(var.clone())
95        }
96    }
97
98    for var in to_remove {
99        map.remove(&var);
100    }
101
102    map
103}
104
105fn hash_lc<E: Engine>(
106    terms: &[(Variable, E::Fr)],
107    h: &mut Blake2s
108)
109{
110    let map = proc_lc::<E>(terms);
111
112    let mut buf = [0u8; 9 + 32];
113    BigEndian::write_u64(&mut buf[0..8], map.len() as u64);
114    h.update(&buf[0..8]);
115
116    for (var, coeff) in map {
117        match var.0.get_unchecked() {
118            Index::Input(i) => {
119                buf[0] = b'I';
120                BigEndian::write_u64(&mut buf[1..9], i as u64);
121            },
122            Index::Aux(i) => {
123                buf[0] = b'A';
124                BigEndian::write_u64(&mut buf[1..9], i as u64);
125            }
126        }
127        
128        coeff.into_repr().write_be(&mut buf[9..]).unwrap();
129
130        h.update(&buf);
131    }
132}
133
134fn eval_lc<E: Engine>(
135    terms: &[(Variable, E::Fr)],
136    inputs: &[(E::Fr, String)],
137    aux: &[(E::Fr, String)]
138) -> E::Fr
139{
140    let mut acc = E::Fr::zero();
141
142    for &(var, ref coeff) in terms {
143        let mut tmp = match var.get_unchecked() {
144            Index::Input(index) => inputs[index].0,
145            Index::Aux(index) => aux[index].0
146        };
147
148        tmp.mul_assign(&coeff);
149        acc.add_assign(&tmp);
150    }
151
152    acc
153}
154
155impl<E: Engine> TestConstraintSystem<E> {
156    pub fn new() -> TestConstraintSystem<E> {
157        let mut map = HashMap::new();
158        map.insert("ONE".into(), NamedObject::Var(TestConstraintSystem::<E>::one()));
159
160        TestConstraintSystem {
161            named_objects: map,
162            current_namespace: vec![],
163            constraints: vec![],
164            inputs: vec![(E::Fr::one(), "ONE".into())],
165            aux: vec![]
166        }
167    }
168
169    pub fn pretty_print(&self) -> String {
170        let mut s = String::new();
171
172        let negone = {
173            let mut tmp = E::Fr::one();
174            tmp.negate();
175            tmp
176        };
177
178        let powers_of_two = (0..E::Fr::NUM_BITS).map(|i| {
179            E::Fr::from_str("2").unwrap().pow(&[i as u64])
180        }).collect::<Vec<_>>();
181
182        let pp = |s: &mut String, lc: &LinearCombination<E>| {
183            write!(s, "(").unwrap();
184            let mut is_first = true;
185            for (var, coeff) in proc_lc::<E>(lc.as_ref()) {
186                if coeff == negone {
187                    write!(s, " - ").unwrap();
188                } else if !is_first {
189                    write!(s, " + ").unwrap();
190                }
191                is_first = false;
192
193                if coeff != E::Fr::one() && coeff != negone {
194                    for (i, x) in powers_of_two.iter().enumerate() {
195                        if x == &coeff {
196                            write!(s, "2^{} . ", i).unwrap();
197                            break;
198                        }
199                    }
200
201                    write!(s, "{} . ", coeff).unwrap();
202                }
203
204                match var.0.get_unchecked() {
205                    Index::Input(i) => {
206                        write!(s, "`{}`", &self.inputs[i].1).unwrap();
207                    },
208                    Index::Aux(i) => {
209                        write!(s, "`{}`", &self.aux[i].1).unwrap();
210                    }
211                }
212            }
213            if is_first {
214                // Nothing was visited, print 0.
215                write!(s, "0").unwrap();
216            }
217            write!(s, ")").unwrap();
218        };
219
220        for &(ref a, ref b, ref c, ref name) in &self.constraints {
221            write!(&mut s, "\n").unwrap();
222
223            write!(&mut s, "{}: ", name).unwrap();
224            pp(&mut s, a);
225            write!(&mut s, " * ").unwrap();
226            pp(&mut s, b);
227            write!(&mut s, " = ").unwrap();
228            pp(&mut s, c);
229        }
230
231        write!(&mut s, "\n").unwrap();
232
233        s
234    }
235
236    pub fn find_unconstrained(&self) -> String {
237        let mut s = String::new();
238        let pp = |hm: & mut HashSet<String>, lc: &LinearCombination<E>| {
239            for (var, coeff) in proc_lc::<E>(lc.as_ref()) {
240                match var.0.get_unchecked() {
241                    Index::Input(i) => {
242                        let v = self.inputs[i].clone();
243                        hm.insert(v.1);
244                    },
245                    Index::Aux(i) => {
246                        let v = self.aux[i].clone();
247                        hm.insert(v.1);
248                    }
249                }
250            }
251        };
252
253        let i_max = self.constraints.len();
254
255        let mut set = HashSet::new();
256        for &(ref a, ref b, ref c, ref _name) in &self.constraints {
257
258            pp(&mut set, a);
259            pp(&mut set, b);
260            pp(&mut set, c);
261        }
262
263        for inp in self.inputs.iter() {
264            if set.get(&inp.1).is_none() {
265                write!(&mut s, "\n").unwrap();
266                write!(&mut s, "{}", inp.1).unwrap();
267                write!(&mut s, "\n").unwrap();
268            }
269        }
270
271        for inp in self.aux.iter() {
272            if set.get(&inp.1).is_none() {
273                write!(&mut s, "\n").unwrap();
274                write!(&mut s, "{}", inp.1).unwrap();
275                write!(&mut s, "\n").unwrap();
276            }
277        }
278
279        s
280    }
281
282    pub fn hash(&self) -> String {
283        let mut h = Blake2s::new(32);
284        {
285            let mut buf = [0u8; 24];
286
287            BigEndian::write_u64(&mut buf[0..8], self.inputs.len() as u64);
288            BigEndian::write_u64(&mut buf[8..16], self.aux.len() as u64);
289            BigEndian::write_u64(&mut buf[16..24], self.constraints.len() as u64);
290            h.update(&buf);
291        }
292
293        for constraint in &self.constraints {
294            hash_lc::<E>(constraint.0.as_ref(), &mut h);
295            hash_lc::<E>(constraint.1.as_ref(), &mut h);
296            hash_lc::<E>(constraint.2.as_ref(), &mut h);
297        }
298
299        let mut s = String::new();
300        for b in h.finalize().as_ref() {
301            s += &format!("{:02x}", b);
302        }
303
304        s
305    }
306
307    pub fn which_is_unsatisfied(&self) -> Option<&str> {
308        for &(ref a, ref b, ref c, ref path) in &self.constraints {
309            let mut a = eval_lc::<E>(a.as_ref(), &self.inputs, &self.aux);
310            let b = eval_lc::<E>(b.as_ref(), &self.inputs, &self.aux);
311            let c = eval_lc::<E>(c.as_ref(), &self.inputs, &self.aux);
312
313            a.mul_assign(&b);
314
315            if a != c {
316                return Some(&*path)
317            }
318        }
319
320        None
321    }
322
323    pub fn is_satisfied(&self) -> bool
324    {
325        self.which_is_unsatisfied().is_none()
326    }
327
328    pub fn num_constraints(&self) -> usize
329    {
330        self.constraints.len()
331    }
332
333    pub fn set(&mut self, path: &str, to: E::Fr)
334    {
335        match self.named_objects.get(path) {
336            Some(&NamedObject::Var(ref v)) => {
337                match v.get_unchecked() {
338                    Index::Input(index) => self.inputs[index].0 = to,
339                    Index::Aux(index) => self.aux[index].0 = to
340                }
341            }
342            Some(e) => panic!("tried to set path `{}` to value, but `{:?}` already exists there.", path, e),
343            _ => panic!("no variable exists at path: {}", path)
344        }
345    }
346
347    pub fn verify(&self, expected: &[E::Fr]) -> bool
348    {
349        assert_eq!(expected.len() + 1, self.inputs.len());
350
351        for (a, b) in self.inputs.iter().skip(1).zip(expected.iter())
352        {
353            if &a.0 != b {
354                return false
355            }
356        }
357
358        return true;
359    }
360
361    pub fn num_inputs(&self) -> usize {
362        self.inputs.len()
363    }
364
365    pub fn get_input(&mut self, index: usize, path: &str) -> E::Fr
366    {
367        let (assignment, name) = self.inputs[index].clone();
368
369        assert_eq!(path, name);
370
371        assignment
372    }
373
374    pub fn get(&mut self, path: &str) -> E::Fr
375    {
376        match self.named_objects.get(path) {
377            Some(&NamedObject::Var(ref v)) => {
378                match v.get_unchecked() {
379                    Index::Input(index) => self.inputs[index].0,
380                    Index::Aux(index) => self.aux[index].0
381                }
382            }
383            Some(e) => panic!("tried to get value of path `{}`, but `{:?}` exists there (not a variable)", path, e),
384            _ => panic!("no variable exists at path: {}", path)
385        }
386    }
387
388    fn set_named_obj(&mut self, path: String, to: NamedObject) {
389        if self.named_objects.contains_key(&path) {
390            panic!("tried to create object at existing path: {}", path);
391        }
392
393        self.named_objects.insert(path, to);
394    }
395}
396
397fn compute_path(ns: &[String], this: String) -> String {
398    if this.chars().any(|a| a == '/') {
399        panic!("'/' is not allowed in names");
400    }
401
402    let mut name = String::new();
403
404    let mut needs_separation = false;
405    for ns in ns.iter().chain(Some(&this).into_iter())
406    {
407        if needs_separation {
408            name += "/";
409        }
410
411        name += ns;
412        needs_separation = true;
413    }
414
415    name
416}
417
418impl<E: Engine> ConstraintSystem<E> for TestConstraintSystem<E> {
419    type Root = Self;
420
421    fn alloc<F, A, AR>(
422        &mut self,
423        annotation: A,
424        f: F
425    ) -> Result<Variable, SynthesisError>
426        where F: FnOnce() -> Result<E::Fr, SynthesisError>, A: FnOnce() -> AR, AR: Into<String>
427    {
428        let index = self.aux.len();
429        let path = compute_path(&self.current_namespace, annotation().into());
430        self.aux.push((f()?, path.clone()));
431        let var = Variable::new_unchecked(Index::Aux(index));
432        self.set_named_obj(path, NamedObject::Var(var));
433
434        Ok(var)
435    }
436
437    fn alloc_input<F, A, AR>(
438        &mut self,
439        annotation: A,
440        f: F
441    ) -> Result<Variable, SynthesisError>
442        where F: FnOnce() -> Result<E::Fr, SynthesisError>, A: FnOnce() -> AR, AR: Into<String>
443    {
444        let index = self.inputs.len();
445        let path = compute_path(&self.current_namespace, annotation().into());
446        self.inputs.push((f()?, path.clone()));
447        let var = Variable::new_unchecked(Index::Input(index));
448        self.set_named_obj(path, NamedObject::Var(var));
449
450        Ok(var)
451    }
452
453    fn enforce<A, AR, LA, LB, LC>(
454        &mut self,
455        annotation: A,
456        a: LA,
457        b: LB,
458        c: LC
459    )
460        where A: FnOnce() -> AR, AR: Into<String>,
461              LA: FnOnce(LinearCombination<E>) -> LinearCombination<E>,
462              LB: FnOnce(LinearCombination<E>) -> LinearCombination<E>,
463              LC: FnOnce(LinearCombination<E>) -> LinearCombination<E>
464    {
465        let path = compute_path(&self.current_namespace, annotation().into());
466        let index = self.constraints.len();
467        self.set_named_obj(path.clone(), NamedObject::Constraint(index));
468
469        let a = a(LinearCombination::zero());
470        let b = b(LinearCombination::zero());
471        let c = c(LinearCombination::zero());
472
473        self.constraints.push((a, b, c, path));
474    }
475
476    fn push_namespace<NR, N>(&mut self, name_fn: N)
477    where NR: Into<String>, N: FnOnce() -> NR
478    {
479        let name = name_fn().into();
480        let path = compute_path(&self.current_namespace, name.clone());
481        self.set_named_obj(path.clone(), NamedObject::Namespace);
482        self.current_namespace.push(name);
483    }
484
485    fn pop_namespace(&mut self)
486    {
487        assert!(self.current_namespace.pop().is_some());
488    }
489
490    fn get_root(&mut self) -> &mut Self::Root
491    {
492        self
493    }
494}
495
496#[test]
497fn test_cs() {
498    use bellman::pairing::bls12_381::{Bls12, Fr};
499    use bellman::pairing::ff::PrimeField;
500
501    let mut cs = TestConstraintSystem::<Bls12>::new();
502    assert!(cs.is_satisfied());
503    assert_eq!(cs.num_constraints(), 0);
504    let a = cs.namespace(|| "a").alloc(|| "var", || Ok(Fr::from_str("10").unwrap())).unwrap();
505    let b = cs.namespace(|| "b").alloc(|| "var", || Ok(Fr::from_str("4").unwrap())).unwrap();
506    let c = cs.alloc(|| "product", || Ok(Fr::from_str("40").unwrap())).unwrap();
507
508    cs.enforce(
509        || "mult",
510        |lc| lc + a,
511        |lc| lc + b,
512        |lc| lc + c
513    );
514    assert!(cs.is_satisfied());
515    assert_eq!(cs.num_constraints(), 1);
516
517    cs.set("a/var", Fr::from_str("4").unwrap());
518
519    let one = TestConstraintSystem::<Bls12>::one();
520    cs.enforce(
521        || "eq",
522        |lc| lc + a,
523        |lc| lc + one,
524        |lc| lc + b
525    );
526
527    assert!(!cs.is_satisfied());
528    assert!(cs.which_is_unsatisfied() == Some("mult"));
529
530    assert!(cs.get("product") == Fr::from_str("40").unwrap());
531
532    cs.set("product", Fr::from_str("16").unwrap());
533    assert!(cs.is_satisfied());
534
535    {
536        let mut cs = cs.namespace(|| "test1");
537        let mut cs = cs.namespace(|| "test2");
538        cs.alloc(|| "hehe", || Ok(Fr::one())).unwrap();
539    }
540
541    assert!(cs.get("test1/test2/hehe") == Fr::one());
542}