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 crate::utils::boolean_hypercube;
260
261 use super::*;
262 type F = halo2curves::secp256k1::Fp;
263 use crate::polynomial::ml_poly::MlPoly;
264
265 #[test]
266 fn test_r1cs() {
267 let num_cons = 10;
268 let num_input = 3;
269 let num_vars = num_cons - num_input;
270
271 let (r1cs, mut witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
272
273 assert_eq!(witness.len(), num_vars);
274 assert_eq!(r1cs.public_input.len(), num_input);
275
276 assert!(r1cs.is_sat(&witness, &r1cs.public_input));
277
278 // Should assert if the witness is invalid
279 witness[0] = witness[0] + F::ONE;
280 assert!(r1cs.is_sat(&witness, &r1cs.public_input) == false);
281 witness[0] = witness[0] - F::ONE;
282
283 // Should assert if the public input is invalid
284 let mut public_input = r1cs.public_input.clone();
285 public_input[0] = public_input[0] + F::ONE;
286 assert!(r1cs.is_sat(&witness, &public_input) == false);
287 public_input[0] = public_input[0] - F::ONE;
288
289 // Test MLE
290 let A_mle = r1cs.A.to_ml_extension();
291 let B_mle = r1cs.B.to_ml_extension();
292 let C_mle = r1cs.C.to_ml_extension();
293 let z = R1CS::construct_z(&witness, &public_input);
294 let Z_mle = MlPoly::new(z);
295
296 let s = Z_mle.num_vars;
297 for c in &boolean_hypercube(s) {
298 let mut eval_a = F::ZERO;
299 let mut eval_b = F::ZERO;
300 let mut eval_c = F::ZERO;
301 for b in &boolean_hypercube(s) {
302 let z_eval = Z_mle.eval(&b);
303 let eval_matrix = [c.as_slice(), b.as_slice()].concat();
304 eval_a += A_mle.eval(&eval_matrix) * z_eval;
305 eval_b += B_mle.eval(&eval_matrix) * z_eval;
306 eval_c += C_mle.eval(&eval_matrix) * z_eval;
307 }
308 let eval_con = eval_a * eval_b - eval_c;
309 assert_eq!(eval_con, F::ZERO);
310 }
311 }
312
313 #[test]
314 fn test_construct_z() {
315 let num_cons = 10;
316 let num_input = 3;
317 let num_vars = num_cons - num_input;
318
319 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
320
321 let Z = R1CS::construct_z(&witness, &r1cs.public_input);
322 // The first num_vars should equal to Z
323 let Z_mle = SparseMLPoly::from_dense(Z.clone());
324
325 for (i, b) in boolean_hypercube(Z_mle.num_vars - 1).iter().enumerate() {
326 assert_eq!(Z[i], Z_mle.eval(&[&[F::ZERO], b.as_slice()].concat()));
327 }
328
329 for (i, b) in boolean_hypercube(Z_mle.num_vars - 1).iter().enumerate() {
330 if i == 0 {
331 assert_eq!(F::ONE, Z_mle.eval(&[&[F::ONE], b.as_slice()].concat()));
332 } else if (i - 1) < r1cs.public_input.len() {
333 assert_eq!(
334 r1cs.public_input[i - 1],
335 Z_mle.eval(&[&[F::ONE], b.as_slice()].concat())
336 );
337 } else {
338 assert_eq!(F::ZERO, Z_mle.eval(&[&[F::ONE], b.as_slice()].concat()));
339 }
340 }
341 }
342
343 #[test]
344 fn test_z_len() {
345 let num_cons = 10;
346 let num_input = 3;
347 let num_vars = num_cons - num_input;
348
349 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
350
351 let z = R1CS::construct_z(&witness, &r1cs.public_input);
352 assert_eq!(z.len(), r1cs.z_len());
353 }
354}