Skip to main content

bellman/gadgets/test/
mod.rs

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