shockwave-plus/shockwave_plus/src/r1cs/r1cs.rs
shockwave_plus/src/r1cs/r1cs.rsRaw
1use crate::polynomial::sparse_ml_poly::SparseMLPoly;
2use crate::FieldExt;
3
4#[derive(Clone, Debug)]
5pub struct SparseMatrixEntry<F: FieldExt> {
6 pub row: usize,
7 pub col: usize,
8 pub val: F,
9}
10
11#[derive(Clone)]
12pub struct Matrix<F: FieldExt> {
13 pub entries: Vec<SparseMatrixEntry<F>>,
14 pub num_cols: usize,
15 pub num_rows: usize,
16}
17
18impl<F> Matrix<F>
19where
20 F: FieldExt,
21{
22 pub fn new(entries: Vec<SparseMatrixEntry<F>>, num_cols: usize, num_rows: usize) -> Self {
23 Self {
24 entries,
25 num_cols,
26 num_rows,
27 }
28 }
29
30 pub fn mul_vector(&self, vec: &[F]) -> Vec<F> {
31 debug_assert_eq!(vec.len(), self.num_cols);
32 let mut result = vec![F::ZERO; self.num_rows];
33 let entries = &self.entries;
34 for i in 0..entries.len() {
35 let row = entries[i].row;
36 let col = entries[i].col;
37 let val = entries[i].val;
38 result[row] += val * vec[col];
39 }
40 result
41 }
42
43 // Return a multilinear extension of the matrix
44 // with num_vars * num_vars entries
45 pub fn to_ml_extension(&self) -> SparseMLPoly<F> {
46 let mut evals = Vec::with_capacity(self.entries.len());
47 let entries = &self.entries;
48 let num_cols = self.num_cols;
49 for i in 0..entries.len() {
50 let row = entries[i].row;
51 let col = entries[i].col;
52 let val = entries[i].val;
53 evals.push(((row * num_cols) + col, val));
54 }
55 let ml_poly_num_vars = ((self.num_cols.next_power_of_two()
56 * self.num_rows.next_power_of_two()) as f64)
57 .log2() as usize;
58
59 let ml_poly = SparseMLPoly::new(evals, ml_poly_num_vars);
60 ml_poly
61 }
62
63 /*
64 pub fn fast_to_coeffs(&self, s: usize, x: F) -> Vec<F> {
65 let mut result = F::ZERO;
66 for entry in &self.0 {
67 let row = entry.0;
68 let col = entry.1;
69 let val = entry.2;
70
71 let index = row * 2usize.pow(s as u32) + col;
72 // Get the degrees of the nonzero coefficients
73 // Tensor product (1 - x_0)(1 - x_1)
74 let base = index;
75 let zero_bits = degree & !base;
76
77 let mut zero_bit_degrees = vec![];
78 for j in 0..s {
79 if zero_bits & (1 << j) != 0 {
80 zero_bit_degrees.push(j);
81 }
82 }
83
84 let mut term = val;
85 for degree in zero_bit_degrees {
86 term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
87 }
88 result += term;
89 }
90 result
91 }
92 */
93
94 /*
95 pub fn fast_uni_eval(&self, s: usize, x: F) -> F {
96 let degree = 2usize.pow(s as u32);
97
98 let mut result = F::ZERO;
99 for entry in &self.0 {
100 let row = entry.0;
101 let col = entry.1;
102 let val = entry.2;
103
104 let index = row * 2usize.pow(s as u32) + col;
105 // Get the degrees of the nonzero coefficients
106 // Tensor product (1 - x_0)(1 - x_1)
107 let base = index;
108 let zero_bits = degree & !base;
109
110 let mut zero_bit_degrees = vec![];
111 for j in 0..s {
112 if zero_bits & (1 << j) != 0 {
113 zero_bit_degrees.push(j);
114 }
115 }
116
117 let mut term = val;
118 for degree in zero_bit_degrees {
119 term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
120 }
121 result += term;
122 }
123 result
124 }
125 */
126}
127
128#[derive(Clone)]
129pub struct R1CS<F>
130where
131 F: FieldExt,
132{
133 pub A: Matrix<F>,
134 pub B: Matrix<F>,
135 pub C: Matrix<F>,
136 pub public_input: Vec<F>,
137 pub num_vars: usize,
138 pub num_input: usize,
139}
140
141impl<F> R1CS<F>
142where
143 F: FieldExt,
144{
145 pub fn hadamard_prod(a: &[F], b: &[F]) -> Vec<F> {
146 assert_eq!(a.len(), b.len());
147 let mut result = vec![F::ZERO; a.len()];
148 for i in 0..a.len() {
149 result[i] = a[i] * b[i];
150 }
151 result
152 }
153
154 pub fn num_cons(&self) -> usize {
155 self.A.entries.len()
156 }
157
158 pub fn z_len(&self) -> usize {
159 ((self.num_vars.next_power_of_two() + 1) + self.num_input).next_power_of_two()
160 }
161
162 pub fn construct_z(witness: &[F], public_input: &[F]) -> Vec<F> {
163 // Z = (witness, 1, io)
164
165 let mut z = witness.to_vec();
166 // Pad the witness part of z to have a power of two length
167 z.resize(z.len().next_power_of_two(), F::ZERO);
168 z.push(F::ONE);
169 z.extend(public_input.clone());
170 // Pad the (1, io) part of z to have a power of two length
171 z.resize(z.len().next_power_of_two(), F::ZERO);
172
173 z
174 }
175
176 pub fn produce_synthetic_r1cs(num_vars: usize, num_input: usize) -> (Self, Vec<F>) {
177 let mut public_input = Vec::with_capacity(num_input);
178 let mut witness = Vec::with_capacity(num_vars);
179
180 for i in 0..num_input {
181 public_input.push(F::from((i + 1) as u64));
182 }
183
184 for i in 0..num_vars {
185 witness.push(F::from((i + 1) as u64));
186 }
187
188 let z = Self::construct_z(&witness, &public_input);
189
190 let mut A_entries: Vec<SparseMatrixEntry<F>> = vec![];
191 let mut B_entries: Vec<SparseMatrixEntry<F>> = vec![];
192 let mut C_entries: Vec<SparseMatrixEntry<F>> = vec![];
193
194 // Constrain the variables
195 for i in 0..num_vars {
196 let A_col = i % num_vars;
197 let B_col = (i + 1) % num_vars;
198 let C_col = (i + 2) % num_vars;
199
200 // For the i'th constraint,
201 // add the value 1 at the (i % num_vars)th column of A, B.
202 // Compute the corresponding C_column value so that A_i * B_i = C_i
203 // we apply multiplication since the Hadamard product is computed for Az ・ Bz,
204
205 // We only _enable_ a single variable in each constraint.
206 let AB = if z[C_col] == F::ZERO { F::ZERO } else { F::ONE };
207
208 A_entries.push(SparseMatrixEntry {
209 row: i,
210 col: A_col,
211 val: AB,
212 });
213 B_entries.push(SparseMatrixEntry {
214 row: i,
215 col: B_col,
216 val: AB,
217 });
218 C_entries.push(SparseMatrixEntry {
219 row: i,
220 col: C_col,
221 val: if z[C_col] == F::ZERO {
222 F::ZERO
223 } else {
224 (z[A_col] * z[B_col]) * z[C_col].invert().unwrap()
225 },
226 });
227 }
228
229 // Constrain the public inputs
230 let input_index_start = num_vars.next_power_of_two() + 1;
231 for i in input_index_start..(input_index_start + num_input) {
232 let A_col = i;
233 let B_col = (i + 1) % input_index_start + num_input;
234 let C_col = (i + 2) % input_index_start + num_input;
235
236 // For the i'th constraint,
237 // add the value 1 at the (i % num_vars)th column of A, B.
238 // Compute the corresponding C_column value so that A_i * B_i = C_i
239 // we apply multiplication since the Hadamard product is computed for Az ・ Bz,
240
241 // We only _enable_ a single variable in each constraint.
242 let AB = if z[C_col] == F::ZERO { F::ZERO } else { F::ONE };
243
244 A_entries.push(SparseMatrixEntry {
245 row: i,
246 col: A_col,
247 val: AB,
248 });
249 B_entries.push(SparseMatrixEntry {
250 row: i,
251 col: B_col,
252 val: AB,
253 });
254 C_entries.push(SparseMatrixEntry {
255 row: i,
256 col: C_col,
257 val: if z[C_col] == F::ZERO {
258 F::ZERO
259 } else {
260 (z[A_col] * z[B_col]) * z[C_col].invert().unwrap()
261 },
262 });
263 }
264
265 let num_cols = z.len();
266 let num_rows = z.len();
267
268 let A = Matrix::new(A_entries, num_cols, num_rows);
269 let B = Matrix::new(B_entries, num_cols, num_rows);
270 let C = Matrix::new(C_entries, num_cols, num_rows);
271
272 (
273 Self {
274 A,
275 B,
276 C,
277 public_input,
278 num_vars,
279 num_input,
280 },
281 witness,
282 )
283 }
284
285 pub fn is_sat(&self, witness: &[F], public_input: &[F]) -> bool {
286 let z = Self::construct_z(witness, public_input);
287 let Az = self.A.mul_vector(&z);
288 let Bz = self.B.mul_vector(&z);
289 let Cz = self.C.mul_vector(&z);
290
291 Self::hadamard_prod(&Az, &Bz) == Cz
292 }
293}
294
295#[cfg(test)]
296mod tests {
297
298 use halo2curves::ff::Field;
299
300 use super::*;
301 type F = halo2curves::secp256k1::Fp;
302 use tensor_pcs::MlPoly;
303
304 // Returns a vector of vectors of length m, where each vector is a boolean vector (big endian)
305 fn boolean_hypercube<F: FieldExt>(m: usize) -> Vec<Vec<F>> {
306 let n = 2usize.pow(m as u32);
307
308 let mut boolean_hypercube = Vec::<Vec<F>>::with_capacity(n);
309
310 for i in 0..n {
311 let mut tmp = Vec::with_capacity(m);
312 for j in 0..m {
313 let i_b = F::from((i >> j & 1) as u64);
314 tmp.push(i_b);
315 }
316 tmp.reverse();
317 boolean_hypercube.push(tmp);
318 }
319
320 boolean_hypercube
321 }
322
323 #[test]
324 fn test_r1cs() {
325 let num_cons = 10;
326 let num_input = 3;
327 let num_vars = num_cons - num_input;
328
329 let (r1cs, mut witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
330 assert_eq!(r1cs.num_cons(), num_cons);
331
332 assert_eq!(witness.len(), num_vars);
333 assert_eq!(r1cs.public_input.len(), num_input);
334
335 assert!(r1cs.is_sat(&witness, &r1cs.public_input));
336
337 // Should assert if the witness is invalid
338 witness[0] = witness[0] + F::ONE;
339 assert!(r1cs.is_sat(&witness, &r1cs.public_input) == false);
340 witness[0] = witness[0] - F::ONE;
341
342 // Should assert if the public input is invalid
343 let mut public_input = r1cs.public_input.clone();
344 public_input[0] = public_input[0] + F::ONE;
345 assert!(r1cs.is_sat(&witness, &public_input) == false);
346 public_input[0] = public_input[0] - F::ONE;
347
348 // Test MLE
349 let A_mle = r1cs.A.to_ml_extension();
350 let B_mle = r1cs.B.to_ml_extension();
351 let C_mle = r1cs.C.to_ml_extension();
352 let z = R1CS::construct_z(&witness, &public_input);
353 let Z_mle = MlPoly::new(z);
354
355 let s = Z_mle.num_vars;
356 for c in &boolean_hypercube(s) {
357 let mut eval_a = F::ZERO;
358 let mut eval_b = F::ZERO;
359 let mut eval_c = F::ZERO;
360 for b in &boolean_hypercube(s) {
361 let z_eval = Z_mle.eval(&b);
362 let eval_matrix = [c.as_slice(), b.as_slice()].concat();
363 eval_a += A_mle.eval(&eval_matrix) * z_eval;
364 eval_b += B_mle.eval(&eval_matrix) * z_eval;
365 eval_c += C_mle.eval(&eval_matrix) * z_eval;
366 }
367 let eval_con = eval_a * eval_b - eval_c;
368 assert_eq!(eval_con, F::ZERO);
369 }
370 }
371
372 #[test]
373 fn test_construct_z() {
374 let num_cons = 10;
375 let num_input = 3;
376 let num_vars = num_cons - num_input;
377
378 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
379
380 let Z = R1CS::construct_z(&witness, &r1cs.public_input);
381 // The first num_vars should equal to Z
382 let Z_mle = SparseMLPoly::from_dense(Z.clone());
383
384 for (i, b) in boolean_hypercube(Z_mle.num_vars - 1).iter().enumerate() {
385 assert_eq!(Z[i], Z_mle.eval(&[&[F::ZERO], b.as_slice()].concat()));
386 }
387
388 for (i, b) in boolean_hypercube(Z_mle.num_vars - 1).iter().enumerate() {
389 if i == 0 {
390 assert_eq!(F::ONE, Z_mle.eval(&[&[F::ONE], b.as_slice()].concat()));
391 } else if (i - 1) < r1cs.public_input.len() {
392 assert_eq!(
393 r1cs.public_input[i - 1],
394 Z_mle.eval(&[&[F::ONE], b.as_slice()].concat())
395 );
396 } else {
397 assert_eq!(F::ZERO, Z_mle.eval(&[&[F::ONE], b.as_slice()].concat()));
398 }
399 }
400 }
401
402 #[test]
403 fn test_z_len() {
404 let num_cons = 10;
405 let num_input = 3;
406 let num_vars = num_cons - num_input;
407
408 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
409
410 let z = R1CS::construct_z(&witness, &r1cs.public_input);
411 assert_eq!(z.len(), r1cs.z_len());
412 }
413}