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