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