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 assert!((num_cols * num_rows).is_power_of_two());
24 Self {
25 entries,
26 num_cols,
27 num_rows,
28 }
29 }
30
31 pub fn mul_vector(&self, num_rows: usize, vec: &[F]) -> Vec<F> {
32 let mut result = vec![F::ZERO; 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 produce_synthetic_r1cs(
153 num_cons: usize,
154 num_vars: usize,
155 num_input: usize,
156 ) -> (Self, Vec<F>) {
157 // assert_eq!(num_cons, num_vars);
158 let mut public_input = Vec::with_capacity(num_input);
159 let mut witness = Vec::with_capacity(num_vars);
160
161 for i in 0..num_input {
162 public_input.push(F::from((i + 1) as u64));
163 }
164
165 for i in 0..num_vars {
166 witness.push(F::from((i + 1) as u64));
167 }
168
169 let z: Vec<F> = vec![public_input.clone(), witness.clone()].concat();
170
171 let mut A_entries: Vec<SparseMatrixEntry<F>> = vec![];
172 let mut B_entries: Vec<SparseMatrixEntry<F>> = vec![];
173 let mut C_entries: Vec<SparseMatrixEntry<F>> = vec![];
174
175 for i in 0..num_cons {
176 let A_col = i % num_vars;
177 let B_col = (i + 1) % num_vars;
178 let C_col = (i + 2) % num_vars;
179
180 // For the i'th constraint,
181 // add the value 1 at the (i % num_vars)th column of A, B.
182 // Compute the corresponding C_column value so that A_i * B_i = C_i
183 // we apply multiplication since the Hadamard product is computed for Az ・ Bz,
184
185 // We only _enable_ a single variable in each constraint.
186 A_entries.push(SparseMatrixEntry {
187 row: i,
188 col: A_col,
189 val: F::ONE,
190 });
191 B_entries.push(SparseMatrixEntry {
192 row: i,
193 col: B_col,
194 val: F::ONE,
195 });
196 C_entries.push(SparseMatrixEntry {
197 row: i,
198 col: C_col,
199 val: (z[A_col] * z[B_col]) * z[C_col].invert().unwrap(),
200 });
201 }
202
203 let A = Matrix::new(A_entries, num_vars, num_cons);
204
205 let B = Matrix::new(B_entries, num_vars, num_cons);
206
207 let C = Matrix::new(C_entries, num_vars, num_cons);
208
209 (
210 Self {
211 A,
212 B,
213 C,
214 public_input,
215 num_cons,
216 num_vars,
217 num_input,
218 },
219 witness,
220 )
221 }
222
223 pub fn is_sat(&self, witness: &Vec<F>, public_input: &Vec<F>) -> bool {
224 let mut z = Vec::with_capacity(witness.len() + public_input.len() + 1);
225 z.extend(public_input);
226 z.extend(witness);
227
228 let Az = self.A.mul_vector(self.num_cons, &z);
229 let Bz = self.B.mul_vector(self.num_cons, &z);
230 let Cz = self.C.mul_vector(self.num_cons, &z);
231
232 Self::hadamard_prod(&Az, &Bz) == Cz
233 }
234}
235
236#[cfg(test)]
237mod tests {
238 use crate::utils::boolean_hypercube;
239
240 use super::*;
241 type F = halo2curves::secp256k1::Fp;
242 use crate::polynomial::ml_poly::MlPoly;
243
244 #[test]
245 fn test_r1cs() {
246 let num_cons = 2usize.pow(5);
247 let num_vars = num_cons;
248 let num_input = 0;
249
250 let (r1cs, mut witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
251
252 assert_eq!(witness.len(), num_vars);
253 assert_eq!(r1cs.public_input.len(), num_input);
254
255 assert!(r1cs.is_sat(&witness, &r1cs.public_input));
256
257 // Should assert if the witness is invalid
258 witness[0] = witness[0] + F::one();
259 assert!(r1cs.is_sat(&r1cs.public_input, &witness) == false);
260 witness[0] = witness[0] - F::one();
261
262 /*
263 // Should assert if the public input is invalid
264 let mut public_input = r1cs.public_input.clone();
265 public_input[0] = public_input[0] + F::one();
266 assert!(r1cs.is_sat(&witness, &public_input) == false);
267 */
268
269 // Test MLE
270 let s = (num_vars as f64).log2() as usize;
271 let A_mle = r1cs.A.to_ml_extension();
272 let B_mle = r1cs.B.to_ml_extension();
273 let C_mle = r1cs.C.to_ml_extension();
274 let Z_mle = MlPoly::new(witness);
275
276 for c in &boolean_hypercube(s) {
277 let mut eval_a = F::zero();
278 let mut eval_b = F::zero();
279 let mut eval_c = F::zero();
280 for b in &boolean_hypercube(s) {
281 let mut b_rev = b.clone();
282 b_rev.reverse();
283 let z_eval = Z_mle.eval(&b_rev);
284 let mut eval_matrix = [b.as_slice(), c.as_slice()].concat();
285 eval_matrix.reverse();
286 eval_a += A_mle.eval(&eval_matrix) * z_eval;
287 eval_b += B_mle.eval(&eval_matrix) * z_eval;
288 eval_c += C_mle.eval(&eval_matrix) * z_eval;
289 }
290 let eval_con = eval_a * eval_b - eval_c;
291 assert_eq!(eval_con, F::zero());
292 }
293 }
294
295 /*
296 #[test]
297 fn test_fast_uni_eval() {
298 let (r1cs, _) = R1CS::<F>::produce_synthetic_r1cs(8, 8, 0);
299
300 let eval_at = F::from(33);
301 let result = r1cs.A.fast_uni_eval(r1cs.num_vars, eval_at);
302 println!("result: {:?}", result);
303 }
304 */
305}