Nova/src/bellperson/shape_cs.rs
src/bellperson/shape_cs.rsRaw
1//! Support for generating R1CS shape using bellperson.
2
3use std::{
4 cmp::Ordering,
5 collections::{BTreeMap, HashMap},
6};
7
8use crate::traits::Group;
9use bellperson::{ConstraintSystem, Index, LinearCombination, SynthesisError, Variable};
10use core::fmt::Write;
11use ff::{Field, PrimeField};
12
13#[derive(Clone, Copy)]
14struct OrderedVariable(Variable);
15
16#[derive(Debug)]
17enum NamedObject {
18 Constraint(usize),
19 Var(Variable),
20 Namespace,
21}
22
23impl Eq for OrderedVariable {}
24impl PartialEq for OrderedVariable {
25 fn eq(&self, other: &OrderedVariable) -> bool {
26 match (self.0.get_unchecked(), other.0.get_unchecked()) {
27 (Index::Input(ref a), Index::Input(ref b)) => a == b,
28 (Index::Aux(ref a), Index::Aux(ref b)) => a == b,
29 _ => false,
30 }
31 }
32}
33impl PartialOrd for OrderedVariable {
34 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
35 Some(self.cmp(other))
36 }
37}
38impl Ord for OrderedVariable {
39 fn cmp(&self, other: &Self) -> Ordering {
40 match (self.0.get_unchecked(), other.0.get_unchecked()) {
41 (Index::Input(ref a), Index::Input(ref b)) => a.cmp(b),
42 (Index::Aux(ref a), Index::Aux(ref b)) => a.cmp(b),
43 (Index::Input(_), Index::Aux(_)) => Ordering::Less,
44 (Index::Aux(_), Index::Input(_)) => Ordering::Greater,
45 }
46 }
47}
48
49#[allow(clippy::upper_case_acronyms)]
50/// `ShapeCS` is a `ConstraintSystem` for creating `R1CSShape`s for a circuit.
51pub struct ShapeCS<G: Group>
52where
53 G::Scalar: PrimeField + Field,
54{
55 named_objects: HashMap<String, NamedObject>,
56 current_namespace: Vec<String>,
57 #[allow(clippy::type_complexity)]
58 /// All constraints added to the `ShapeCS`.
59 pub constraints: Vec<(
60 LinearCombination<G::Scalar>,
61 LinearCombination<G::Scalar>,
62 LinearCombination<G::Scalar>,
63 String,
64 )>,
65 inputs: Vec<String>,
66 aux: Vec<String>,
67}
68
69fn proc_lc<Scalar: PrimeField>(
70 terms: &LinearCombination<Scalar>,
71) -> BTreeMap<OrderedVariable, Scalar> {
72 let mut map = BTreeMap::new();
73 for (var, &coeff) in terms.iter() {
74 map
75 .entry(OrderedVariable(var))
76 .or_insert_with(|| Scalar::ZERO)
77 .add_assign(&coeff);
78 }
79
80 // Remove terms that have a zero coefficient to normalize
81 let mut to_remove = vec![];
82 for (var, coeff) in map.iter() {
83 if coeff.is_zero().into() {
84 to_remove.push(*var)
85 }
86 }
87
88 for var in to_remove {
89 map.remove(&var);
90 }
91
92 map
93}
94
95impl<G: Group> ShapeCS<G>
96where
97 G::Scalar: PrimeField,
98{
99 /// Create a new, default `ShapeCS`,
100 pub fn new() -> Self {
101 ShapeCS::default()
102 }
103
104 /// Returns the number of constraints defined for this `ShapeCS`.
105 pub fn num_constraints(&self) -> usize {
106 self.constraints.len()
107 }
108
109 /// Returns the number of inputs defined for this `ShapeCS`.
110 pub fn num_inputs(&self) -> usize {
111 self.inputs.len()
112 }
113
114 /// Returns the number of aux inputs defined for this `ShapeCS`.
115 pub fn num_aux(&self) -> usize {
116 self.aux.len()
117 }
118
119 /// Print all public inputs, aux inputs, and constraint names.
120 #[allow(dead_code)]
121 pub fn pretty_print_list(&self) -> Vec<String> {
122 let mut result = Vec::new();
123
124 for input in &self.inputs {
125 result.push(format!("INPUT {input}"));
126 }
127 for aux in &self.aux {
128 result.push(format!("AUX {aux}"));
129 }
130
131 for (_a, _b, _c, name) in &self.constraints {
132 result.push(name.to_string());
133 }
134
135 result
136 }
137
138 /// Print all iputs and a detailed representation of each constraint.
139 #[allow(dead_code)]
140 pub fn pretty_print(&self) -> String {
141 let mut s = String::new();
142
143 for input in &self.inputs {
144 writeln!(s, "INPUT {}", &input).unwrap()
145 }
146
147 let negone = -<G::Scalar>::ONE;
148
149 let powers_of_two = (0..G::Scalar::NUM_BITS)
150 .map(|i| G::Scalar::from(2u64).pow_vartime([u64::from(i)]))
151 .collect::<Vec<_>>();
152
153 let pp = |s: &mut String, lc: &LinearCombination<G::Scalar>| {
154 s.push('(');
155 let mut is_first = true;
156 for (var, coeff) in proc_lc::<G::Scalar>(lc) {
157 if coeff == negone {
158 s.push_str(" - ")
159 } else if !is_first {
160 s.push_str(" + ")
161 }
162 is_first = false;
163
164 if coeff != <G::Scalar>::ONE && coeff != negone {
165 for (i, x) in powers_of_two.iter().enumerate() {
166 if x == &coeff {
167 write!(s, "2^{i} . ").unwrap();
168 break;
169 }
170 }
171
172 write!(s, "{coeff:?} . ").unwrap()
173 }
174
175 match var.0.get_unchecked() {
176 Index::Input(i) => {
177 write!(s, "`I{}`", &self.inputs[i]).unwrap();
178 }
179 Index::Aux(i) => {
180 write!(s, "`A{}`", &self.aux[i]).unwrap();
181 }
182 }
183 }
184 if is_first {
185 // Nothing was visited, print 0.
186 s.push('0');
187 }
188 s.push(')');
189 };
190
191 for (a, b, c, name) in &self.constraints {
192 s.push('\n');
193
194 write!(s, "{name}: ").unwrap();
195 pp(&mut s, a);
196 write!(s, " * ").unwrap();
197 pp(&mut s, b);
198 s.push_str(" = ");
199 pp(&mut s, c);
200 }
201
202 s.push('\n');
203
204 s
205 }
206
207 /// Associate `NamedObject` with `path`.
208 /// `path` must not already have an associated object.
209 fn set_named_obj(&mut self, path: String, to: NamedObject) {
210 assert!(
211 !self.named_objects.contains_key(&path),
212 "tried to create object at existing path: {path}"
213 );
214
215 self.named_objects.insert(path, to);
216 }
217}
218
219impl<G: Group> Default for ShapeCS<G>
220where
221 G::Scalar: PrimeField,
222{
223 fn default() -> Self {
224 let mut map = HashMap::new();
225 map.insert("ONE".into(), NamedObject::Var(ShapeCS::<G>::one()));
226 ShapeCS {
227 named_objects: map,
228 current_namespace: vec![],
229 constraints: vec![],
230 inputs: vec![String::from("ONE")],
231 aux: vec![],
232 }
233 }
234}
235
236impl<G: Group> ConstraintSystem<G::Scalar> for ShapeCS<G>
237where
238 G::Scalar: PrimeField,
239{
240 type Root = Self;
241
242 fn alloc<F, A, AR>(&mut self, annotation: A, _f: F) -> Result<Variable, SynthesisError>
243 where
244 F: FnOnce() -> Result<G::Scalar, SynthesisError>,
245 A: FnOnce() -> AR,
246 AR: Into<String>,
247 {
248 let path = compute_path(&self.current_namespace, &annotation().into());
249 self.aux.push(path);
250
251 Ok(Variable::new_unchecked(Index::Aux(self.aux.len() - 1)))
252 }
253
254 fn alloc_input<F, A, AR>(&mut self, annotation: A, _f: F) -> Result<Variable, SynthesisError>
255 where
256 F: FnOnce() -> Result<G::Scalar, SynthesisError>,
257 A: FnOnce() -> AR,
258 AR: Into<String>,
259 {
260 let path = compute_path(&self.current_namespace, &annotation().into());
261 self.inputs.push(path);
262
263 Ok(Variable::new_unchecked(Index::Input(self.inputs.len() - 1)))
264 }
265
266 fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
267 where
268 A: FnOnce() -> AR,
269 AR: Into<String>,
270 LA: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
271 LB: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
272 LC: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
273 {
274 let path = compute_path(&self.current_namespace, &annotation().into());
275 let index = self.constraints.len();
276 self.set_named_obj(path.clone(), NamedObject::Constraint(index));
277
278 let a = a(LinearCombination::zero());
279 let b = b(LinearCombination::zero());
280 let c = c(LinearCombination::zero());
281
282 self.constraints.push((a, b, c, path));
283 }
284
285 fn push_namespace<NR, N>(&mut self, name_fn: N)
286 where
287 NR: Into<String>,
288 N: FnOnce() -> NR,
289 {
290 let name = name_fn().into();
291 let path = compute_path(&self.current_namespace, &name);
292 self.set_named_obj(path, NamedObject::Namespace);
293 self.current_namespace.push(name);
294 }
295
296 fn pop_namespace(&mut self) {
297 assert!(self.current_namespace.pop().is_some());
298 }
299
300 fn get_root(&mut self) -> &mut Self::Root {
301 self
302 }
303}
304
305fn compute_path(ns: &[String], this: &str) -> String {
306 assert!(
307 !this.chars().any(|a| a == '/'),
308 "'/' is not allowed in names"
309 );
310
311 let mut name = String::new();
312
313 let mut needs_separation = false;
314 for ns in ns.iter().chain(Some(this.to_string()).iter()) {
315 if needs_separation {
316 name += "/";
317 }
318
319 name += ns;
320 needs_separation = true;
321 }
322
323 name
324}