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
36pub 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 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 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}