1use 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
32pub 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 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 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 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}