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> {
52 named_objects: HashMap<String, NamedObject>,
53 current_namespace: Vec<String>,
54 #[allow(clippy::type_complexity)]
55 /// All constraints added to the `ShapeCS`.
56 pub constraints: Vec<(
57 LinearCombination<G::Scalar>,
58 LinearCombination<G::Scalar>,
59 LinearCombination<G::Scalar>,
60 String,
61 )>,
62 inputs: Vec<String>,
63 aux: Vec<String>,
64}
65
66fn proc_lc<Scalar: PrimeField>(
67 terms: &LinearCombination<Scalar>,
68) -> BTreeMap<OrderedVariable, Scalar> {
69 let mut map = BTreeMap::new();
70 for (var, &coeff) in terms.iter() {
71 map
72 .entry(OrderedVariable(var))
73 .or_insert_with(|| Scalar::ZERO)
74 .add_assign(&coeff);
75 }
76
77 // Remove terms that have a zero coefficient to normalize
78 let mut to_remove = vec![];
79 for (var, coeff) in map.iter() {
80 if coeff.is_zero().into() {
81 to_remove.push(*var)
82 }
83 }
84
85 for var in to_remove {
86 map.remove(&var);
87 }
88
89 map
90}
91
92impl<G: Group> ShapeCS<G> {
93 /// Create a new, default `ShapeCS`,
94 pub fn new() -> Self {
95 ShapeCS::default()
96 }
97
98 /// Returns the number of constraints defined for this `ShapeCS`.
99 pub fn num_constraints(&self) -> usize {
100 self.constraints.len()
101 }
102
103 /// Returns the number of inputs defined for this `ShapeCS`.
104 pub fn num_inputs(&self) -> usize {
105 self.inputs.len()
106 }
107
108 /// Returns the number of aux inputs defined for this `ShapeCS`.
109 pub fn num_aux(&self) -> usize {
110 self.aux.len()
111 }
112
113 /// Print all public inputs, aux inputs, and constraint names.
114 #[allow(dead_code)]
115 pub fn pretty_print_list(&self) -> Vec<String> {
116 let mut result = Vec::new();
117
118 for input in &self.inputs {
119 result.push(format!("INPUT {input}"));
120 }
121 for aux in &self.aux {
122 result.push(format!("AUX {aux}"));
123 }
124
125 for (_a, _b, _c, name) in &self.constraints {
126 result.push(name.to_string());
127 }
128
129 result
130 }
131
132 /// Print all iputs and a detailed representation of each constraint.
133 #[allow(dead_code)]
134 pub fn pretty_print(&self) -> String {
135 let mut s = String::new();
136
137 for input in &self.inputs {
138 writeln!(s, "INPUT {}", &input).unwrap()
139 }
140
141 let negone = -<G::Scalar>::ONE;
142
143 let powers_of_two = (0..G::Scalar::NUM_BITS)
144 .map(|i| G::Scalar::from(2u64).pow_vartime([u64::from(i)]))
145 .collect::<Vec<_>>();
146
147 let pp = |s: &mut String, lc: &LinearCombination<G::Scalar>| {
148 s.push('(');
149 let mut is_first = true;
150 for (var, coeff) in proc_lc::<G::Scalar>(lc) {
151 if coeff == negone {
152 s.push_str(" - ")
153 } else if !is_first {
154 s.push_str(" + ")
155 }
156 is_first = false;
157
158 if coeff != <G::Scalar>::ONE && coeff != negone {
159 for (i, x) in powers_of_two.iter().enumerate() {
160 if x == &coeff {
161 write!(s, "2^{i} . ").unwrap();
162 break;
163 }
164 }
165
166 write!(s, "{coeff:?} . ").unwrap()
167 }
168
169 match var.0.get_unchecked() {
170 Index::Input(i) => {
171 write!(s, "`I{}`", &self.inputs[i]).unwrap();
172 }
173 Index::Aux(i) => {
174 write!(s, "`A{}`", &self.aux[i]).unwrap();
175 }
176 }
177 }
178 if is_first {
179 // Nothing was visited, print 0.
180 s.push('0');
181 }
182 s.push(')');
183 };
184
185 for (a, b, c, name) in &self.constraints {
186 s.push('\n');
187
188 write!(s, "{name}: ").unwrap();
189 pp(&mut s, a);
190 write!(s, " * ").unwrap();
191 pp(&mut s, b);
192 s.push_str(" = ");
193 pp(&mut s, c);
194 }
195
196 s.push('\n');
197
198 s
199 }
200
201 /// Associate `NamedObject` with `path`.
202 /// `path` must not already have an associated object.
203 fn set_named_obj(&mut self, path: String, to: NamedObject) {
204 assert!(
205 !self.named_objects.contains_key(&path),
206 "tried to create object at existing path: {path}"
207 );
208
209 self.named_objects.insert(path, to);
210 }
211}
212
213impl<G: Group> Default for ShapeCS<G> {
214 fn default() -> Self {
215 let mut map = HashMap::new();
216 map.insert("ONE".into(), NamedObject::Var(ShapeCS::<G>::one()));
217 ShapeCS {
218 named_objects: map,
219 current_namespace: vec![],
220 constraints: vec![],
221 inputs: vec![String::from("ONE")],
222 aux: vec![],
223 }
224 }
225}
226
227impl<G: Group> ConstraintSystem<G::Scalar> for ShapeCS<G> {
228 type Root = Self;
229
230 fn alloc<F, A, AR>(&mut self, annotation: A, _f: F) -> Result<Variable, SynthesisError>
231 where
232 F: FnOnce() -> Result<G::Scalar, SynthesisError>,
233 A: FnOnce() -> AR,
234 AR: Into<String>,
235 {
236 let path = compute_path(&self.current_namespace, &annotation().into());
237 self.aux.push(path);
238
239 Ok(Variable::new_unchecked(Index::Aux(self.aux.len() - 1)))
240 }
241
242 fn alloc_input<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.inputs.push(path);
250
251 Ok(Variable::new_unchecked(Index::Input(self.inputs.len() - 1)))
252 }
253
254 fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
255 where
256 A: FnOnce() -> AR,
257 AR: Into<String>,
258 LA: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
259 LB: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
260 LC: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
261 {
262 let path = compute_path(&self.current_namespace, &annotation().into());
263 let index = self.constraints.len();
264 self.set_named_obj(path.clone(), NamedObject::Constraint(index));
265
266 let a = a(LinearCombination::zero());
267 let b = b(LinearCombination::zero());
268 let c = c(LinearCombination::zero());
269
270 self.constraints.push((a, b, c, path));
271 }
272
273 fn push_namespace<NR, N>(&mut self, name_fn: N)
274 where
275 NR: Into<String>,
276 N: FnOnce() -> NR,
277 {
278 let name = name_fn().into();
279 let path = compute_path(&self.current_namespace, &name);
280 self.set_named_obj(path, NamedObject::Namespace);
281 self.current_namespace.push(name);
282 }
283
284 fn pop_namespace(&mut self) {
285 assert!(self.current_namespace.pop().is_some());
286 }
287
288 fn get_root(&mut self) -> &mut Self::Root {
289 self
290 }
291}
292
293fn compute_path(ns: &[String], this: &str) -> String {
294 assert!(
295 !this.chars().any(|a| a == '/'),
296 "'/' is not allowed in names"
297 );
298
299 let mut name = ns.join("/");
300 if !name.is_empty() {
301 name.push('/');
302 }
303
304 name.push_str(this);
305
306 name
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312
313 #[test]
314 fn test_compute_path() {
315 let ns = vec!["path".to_string(), "to".to_string(), "dir".to_string()];
316 let this = "file";
317 assert_eq!(compute_path(&ns, this), "path/to/dir/file");
318
319 let ns = vec!["".to_string(), "".to_string(), "".to_string()];
320 let this = "file";
321 assert_eq!(compute_path(&ns, this), "///file");
322 }
323
324 #[test]
325 #[should_panic(expected = "'/' is not allowed in names")]
326 fn test_compute_path_invalid() {
327 let ns = vec!["path".to_string(), "to".to_string(), "dir".to_string()];
328 let this = "fi/le";
329 compute_path(&ns, this);
330 }
331}