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