shockwave-plus/shockwave_plus/src/r1cs/r1cs.rs
shockwave_plus/src/r1cs/r1cs.rsRaw
1use crate::FieldExt;
2use halo2curves::ff::Field;
3use tensor_pcs::SparseMLPoly;
4
5#[derive(Clone)]
6pub struct SparseMatrixEntry<F: FieldExt> {
7 pub row: usize,
8 pub col: usize,
9 pub val: F,
10}
11
12#[derive(Clone)]
13pub struct Matrix<F: FieldExt> {
14 pub entries: Vec<SparseMatrixEntry<F>>,
15 pub num_cols: usize,
16 pub num_rows: usize,
17}
18
19impl<F> Matrix<F>
20where
21 F: FieldExt,
22{
23 pub fn new(entries: Vec<SparseMatrixEntry<F>>, num_cols: usize, num_rows: usize) -> Self {
24 assert!((num_cols * num_rows).is_power_of_two());
25 Self {
26 entries,
27 num_cols,
28 num_rows,
29 }
30 }
31
32 pub fn mul_vector(&self, num_rows: usize, vec: &[F]) -> Vec<F> {
33 let mut result = vec![F::ZERO; num_rows];
34 let entries = &self.entries;
35 for i in 0..entries.len() {
36 let row = entries[i].row;
37 let col = entries[i].col;
38 let val = entries[i].val;
39 result[row] += val * vec[col];
40 }
41 result
42 }
43
44 // Return a multilinear extension of the matrix
45 // with num_vars * num_vars entries
46 pub fn to_ml_extension(&self) -> SparseMLPoly<F> {
47 let mut evals = Vec::with_capacity(self.entries.len());
48 let entries = &self.entries;
49 let num_cols = self.num_cols;
50 for i in 0..entries.len() {
51 let row = entries[i].row;
52 let col = entries[i].col;
53 let val = entries[i].val;
54 evals.push(((row * num_cols) + col, val));
55 }
56 let ml_poly_num_vars = ((self.num_cols * self.num_rows) as f64).log2() as usize;
57 let ml_poly = SparseMLPoly::new(evals, ml_poly_num_vars);
58 ml_poly
59 }
60
61 /*
62 pub fn fast_to_coeffs(&self, s: usize, x: F) -> Vec<F> {
63 let mut result = F::ZERO;
64 for entry in &self.0 {
65 let row = entry.0;
66 let col = entry.1;
67 let val = entry.2;
68
69 let index = row * 2usize.pow(s as u32) + col;
70 // Get the degrees of the nonzero coefficients
71 // Tensor product (1 - x_0)(1 - x_1)
72 let base = index;
73 let zero_bits = degree & !base;
74
75 let mut zero_bit_degrees = vec![];
76 for j in 0..s {
77 if zero_bits & (1 << j) != 0 {
78 zero_bit_degrees.push(j);
79 }
80 }
81
82 let mut term = val;
83 for degree in zero_bit_degrees {
84 term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
85 }
86 result += term;
87 }
88 result
89 }
90 */
91
92 /*
93 pub fn fast_uni_eval(&self, s: usize, x: F) -> F {
94 let degree = 2usize.pow(s as u32);
95
96 let mut result = F::ZERO;
97 for entry in &self.0 {
98 let row = entry.0;
99 let col = entry.1;
100 let val = entry.2;
101
102 let index = row * 2usize.pow(s as u32) + col;
103 // Get the degrees of the nonzero coefficients
104 // Tensor product (1 - x_0)(1 - x_1)
105 let base = index;
106 let zero_bits = degree & !base;
107
108 let mut zero_bit_degrees = vec![];
109 for j in 0..s {
110 if zero_bits & (1 << j) != 0 {
111 zero_bit_degrees.push(j);
112 }
113 }
114
115 let mut term = val;
116 for degree in zero_bit_degrees {
117 term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
118 }
119 result += term;
120 }
121 result
122 }
123 */
124}
125
126#[derive(Clone)]
127pub struct R1CS<F>
128where
129 F: FieldExt,
130{
131 pub A: Matrix<F>,
132 pub B: Matrix<F>,
133 pub C: Matrix<F>,
134 pub public_input: Vec<F>,
135 pub num_cons: usize,
136 pub num_vars: usize,
137 pub num_input: usize,
138}
139
140impl<F> R1CS<F>
141where
142 F: FieldExt,
143{
144 pub fn hadamard_prod(a: &[F], b: &[F]) -> Vec<F> {
145 assert_eq!(a.len(), b.len());
146 let mut result = vec![F::ZERO; a.len()];
147 for i in 0..a.len() {
148 result[i] = a[i] * b[i];
149 }
150 result
151 }
152
153 pub fn produce_synthetic_r1cs(
154 num_cons: usize,
155 num_vars: usize,
156 num_input: usize,
157 ) -> (Self, Vec<F>) {
158 // assert_eq!(num_cons, num_vars);
159 let mut public_input = Vec::with_capacity(num_input);
160 let mut witness = Vec::with_capacity(num_vars);
161
162 for i in 0..num_input {
163 public_input.push(F::from((i + 1) as u64));
164 }
165
166 for i in 0..num_vars {
167 witness.push(F::from((i + 1) as u64));
168 }
169
170 let z: Vec<F> = vec![public_input.clone(), witness.clone()].concat();
171
172 let mut A_entries: Vec<SparseMatrixEntry<F>> = vec![];
173 let mut B_entries: Vec<SparseMatrixEntry<F>> = vec![];
174 let mut C_entries: Vec<SparseMatrixEntry<F>> = vec![];
175
176 for i in 0..num_cons {
177 let A_col = i % num_vars;
178 let B_col = (i + 1) % num_vars;
179 let C_col = (i + 2) % num_vars;
180
181 // For the i'th constraint,
182 // add the value 1 at the (i % num_vars)th column of A, B.
183 // Compute the corresponding C_column value so that A_i * B_i = C_i
184 // we apply multiplication since the Hadamard product is computed for Az ・ Bz,
185
186 // We only _enable_ a single variable in each constraint.
187 A_entries.push(SparseMatrixEntry {
188 row: i,
189 col: A_col,
190 val: F::ONE,
191 });
192 B_entries.push(SparseMatrixEntry {
193 row: i,
194 col: B_col,
195 val: F::ONE,
196 });
197 C_entries.push(SparseMatrixEntry {
198 row: i,
199 col: C_col,
200 val: (z[A_col] * z[B_col]) * z[C_col].invert().unwrap(),
201 });
202 }
203
204 let A = Matrix::new(A_entries, num_vars, num_cons);
205
206 let B = Matrix::new(B_entries, num_vars, num_cons);
207
208 let C = Matrix::new(C_entries, num_vars, num_cons);
209
210 (
211 Self {
212 A,
213 B,
214 C,
215 public_input,
216 num_cons,
217 num_vars,
218 num_input,
219 },
220 witness,
221 )
222 }
223
224 pub fn is_sat(&self, witness: &Vec<F>, public_input: &Vec<F>) -> bool {
225 let mut z = Vec::with_capacity(witness.len() + public_input.len() + 1);
226 z.extend(public_input);
227 z.extend(witness);
228
229 let Az = self.A.mul_vector(self.num_cons, &z);
230 let Bz = self.B.mul_vector(self.num_cons, &z);
231 let Cz = self.C.mul_vector(self.num_cons, &z);
232
233 Self::hadamard_prod(&Az, &Bz) == Cz
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use crate::utils::boolean_hypercube;
240
241 use super::*;
242 type F = halo2curves::secp256k1::Fp;
243 use crate::polynomial::ml_poly::MlPoly;
244
245 #[test]
246 fn test_r1cs() {
247 let num_cons = 2usize.pow(5);
248 let num_vars = num_cons;
249 let num_input = 0;
250
251 let (r1cs, mut witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
252
253 assert_eq!(witness.len(), num_vars);
254 assert_eq!(r1cs.public_input.len(), num_input);
255
256 assert!(r1cs.is_sat(&witness, &r1cs.public_input));
257
258 // Should assert if the witness is invalid
259 witness[0] = witness[0] + F::one();
260 assert!(r1cs.is_sat(&r1cs.public_input, &witness) == false);
261 witness[0] = witness[0] - F::one();
262
263 /*
264 // Should assert if the public input is invalid
265 let mut public_input = r1cs.public_input.clone();
266 public_input[0] = public_input[0] + F::one();
267 assert!(r1cs.is_sat(&witness, &public_input) == false);
268 */
269
270 // Test MLE
271 let s = (num_vars as f64).log2() as usize;
272 let A_mle = r1cs.A.to_ml_extension();
273 let B_mle = r1cs.B.to_ml_extension();
274 let C_mle = r1cs.C.to_ml_extension();
275 let Z_mle = MlPoly::new(witness);
276
277 for c in &boolean_hypercube(s) {
278 let mut eval_a = F::zero();
279 let mut eval_b = F::zero();
280 let mut eval_c = F::zero();
281 for b in &boolean_hypercube(s) {
282 let mut b_rev = b.clone();
283 b_rev.reverse();
284 let z_eval = Z_mle.eval(&b_rev);
285 let mut eval_matrix = [b.as_slice(), c.as_slice()].concat();
286 eval_matrix.reverse();
287 eval_a += A_mle.eval(&eval_matrix) * z_eval;
288 eval_b += B_mle.eval(&eval_matrix) * z_eval;
289 eval_c += C_mle.eval(&eval_matrix) * z_eval;
290 }
291 let eval_con = eval_a * eval_b - eval_c;
292 assert_eq!(eval_con, F::zero());
293 }
294 }
295
296 /*
297 #[test]
298 fn test_fast_uni_eval() {
299 let (r1cs, _) = R1CS::<F>::produce_synthetic_r1cs(8, 8, 0);
300
301 let eval_at = F::from(33);
302 let result = r1cs.A.fast_uni_eval(r1cs.num_vars, eval_at);
303 println!("result: {:?}", result);
304 }
305 */
306}