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 assert!(
217 !self.named_objects.contains_key(&path),
218 "tried to create object at existing path: {}",
219 path
220 );
221
222 self.named_objects.insert(path, to);
223 }
224}
225
226impl<G: Group> Default for ShapeCS<G>
227where
228 G::Scalar: PrimeField,
229{
230 fn default() -> Self {
231 let mut map = HashMap::new();
232 map.insert("ONE".into(), NamedObject::Var(ShapeCS::<G>::one()));
233 ShapeCS {
234 named_objects: map,
235 current_namespace: vec![],
236 constraints: vec![],
237 inputs: vec![String::from("ONE")],
238 aux: vec![],
239 }
240 }
241}
242
243impl<G: Group> ConstraintSystem<G::Scalar> for ShapeCS<G>
244where
245 G::Scalar: PrimeField,
246{
247 type Root = Self;
248
249 fn alloc<F, A, AR>(&mut self, annotation: A, _f: F) -> Result<Variable, SynthesisError>
250 where
251 F: FnOnce() -> Result<G::Scalar, SynthesisError>,
252 A: FnOnce() -> AR,
253 AR: Into<String>,
254 {
255 let path = compute_path(&self.current_namespace, &annotation().into());
256 self.aux.push(path);
257
258 Ok(Variable::new_unchecked(Index::Aux(self.aux.len() - 1)))
259 }
260
261 fn alloc_input<F, A, AR>(&mut self, annotation: A, _f: F) -> Result<Variable, SynthesisError>
262 where
263 F: FnOnce() -> Result<G::Scalar, SynthesisError>,
264 A: FnOnce() -> AR,
265 AR: Into<String>,
266 {
267 let path = compute_path(&self.current_namespace, &annotation().into());
268 self.inputs.push(path);
269
270 Ok(Variable::new_unchecked(Index::Input(self.inputs.len() - 1)))
271 }
272
273 fn enforce<A, AR, LA, LB, LC>(&mut self, annotation: A, a: LA, b: LB, c: LC)
274 where
275 A: FnOnce() -> AR,
276 AR: Into<String>,
277 LA: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
278 LB: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
279 LC: FnOnce(LinearCombination<G::Scalar>) -> LinearCombination<G::Scalar>,
280 {
281 let path = compute_path(&self.current_namespace, &annotation().into());
282 let index = self.constraints.len();
283 self.set_named_obj(path.clone(), NamedObject::Constraint(index));
284
285 let a = a(LinearCombination::zero());
286 let b = b(LinearCombination::zero());
287 let c = c(LinearCombination::zero());
288
289 self.constraints.push((a, b, c, path));
290 }
291
292 fn push_namespace<NR, N>(&mut self, name_fn: N)
293 where
294 NR: Into<String>,
295 N: FnOnce() -> NR,
296 {
297 let name = name_fn().into();
298 let path = compute_path(&self.current_namespace, &name);
299 self.set_named_obj(path, NamedObject::Namespace);
300 self.current_namespace.push(name);
301 }
302
303 fn pop_namespace(&mut self) {
304 assert!(self.current_namespace.pop().is_some());
305 }
306
307 fn get_root(&mut self) -> &mut Self::Root {
308 self
309 }
310}
311
312fn compute_path(ns: &[String], this: &str) -> String {
313 assert!(
314 !this.chars().any(|a| a == '/'),
315 "'/' is not allowed in names"
316 );
317
318 let mut name = String::new();
319
320 let mut needs_separation = false;
321 for ns in ns.iter().chain(Some(this.to_string()).iter()) {
322 if needs_separation {
323 name += "/";
324 }
325
326 name += ns;
327 needs_separation = true;
328 }
329
330 name
331}