shockwave-plus/shockwave_plus/src/lib.rs
shockwave_plus/src/lib.rsRaw
1#![allow(non_snake_case)]
2mod polynomial;
3mod r1cs;
4mod sumcheck;
5mod utils;
6
7use ark_std::{end_timer, start_timer};
8use serde::{Deserialize, Serialize};
9use sumcheck::{SCPhase1Proof, SCPhase2Proof, SumCheckPhase1, SumCheckPhase2};
10
11// Exports
12pub use r1cs::R1CS;
13pub use tensor_pcs::*;
14
15#[derive(Serialize, Deserialize)]
16pub struct PartialSpartanProof<F: FieldExt> {
17 pub z_comm: [u8; 32],
18 pub sc_proof_1: SCPhase1Proof<F>,
19 pub sc_proof_2: SCPhase2Proof<F>,
20 pub z_eval_proof: TensorMLOpening<F>,
21 pub v_A: F,
22 pub v_B: F,
23 pub v_C: F,
24}
25
26pub struct FullSpartanProof<F: FieldExt> {
27 pub partial_proof: PartialSpartanProof<F>,
28 pub A_eval_proof: TensorMLOpening<F>,
29 pub B_eval_proof: TensorMLOpening<F>,
30 pub C_eval_proof: TensorMLOpening<F>,
31}
32
33pub struct ShockwavePlus<F: FieldExt> {
34 pub r1cs: R1CS<F>,
35 pub pcs_witness: TensorMultilinearPCS<F>,
36 pub pcs_blinder: TensorMultilinearPCS<F>,
37}
38
39impl<F: FieldExt> ShockwavePlus<F> {
40 pub fn new(r1cs: R1CS<F>, l: usize, num_rows: usize) -> Self {
41 let num_cols = r1cs.num_vars / num_rows;
42
43 // Make sure that there are enough columns to run the l queries
44 assert!(num_cols > l);
45
46 let expansion_factor = 2;
47
48 let ecfft_config = rs_config::ecfft::gen_config(num_cols.next_power_of_two());
49
50 let pcs_config = TensorRSMultilinearPCSConfig::<F> {
51 expansion_factor,
52 domain_powers: None,
53 fft_domain: None,
54 ecfft_config: Some(ecfft_config),
55 l,
56 num_entries: r1cs.num_vars,
57 num_rows,
58 };
59
60 let pcs_witness = TensorMultilinearPCS::new(pcs_config);
61
62 let ecfft_config_blinder =
63 rs_config::ecfft::gen_config((r1cs.z_len() / num_rows).next_power_of_two());
64 let pcs_blinder_config = TensorRSMultilinearPCSConfig::<F> {
65 expansion_factor,
66 domain_powers: None,
67 fft_domain: None,
68 ecfft_config: Some(ecfft_config_blinder),
69 l,
70 num_entries: r1cs.z_len(),
71 num_rows,
72 };
73
74 let pcs_blinder = TensorMultilinearPCS::new(pcs_blinder_config);
75
76 Self {
77 r1cs,
78 pcs_witness,
79 pcs_blinder,
80 }
81 }
82
83 pub fn prove(
84 &self,
85 r1cs_witness: &[F],
86 r1cs_input: &[F],
87 transcript: &mut Transcript<F>,
88 ) -> (PartialSpartanProof<F>, Vec<F>) {
89 // Compute the multilinear extension of the witness
90 let witness_poly = SparseMLPoly::from_dense(r1cs_witness.to_vec());
91 let Z = R1CS::construct_z(r1cs_witness, r1cs_input);
92
93 // Commit the witness polynomial
94 let comm_witness_timer = start_timer!(|| "Commit witness");
95 let committed_witness = self.pcs_witness.commit(&witness_poly);
96 let witness_comm = committed_witness.committed_tree.root;
97 end_timer!(comm_witness_timer);
98
99 // Add the witness commitment to the transcript
100 transcript.append_bytes(&witness_comm);
101
102 // ############################
103 // Phase 1
104 // ###################
105
106 let m = (self.r1cs.z_len() as f64).log2() as usize;
107 let tau = transcript.challenge_vec(m);
108
109 let Az_poly = self.r1cs.A.mul_vector(&Z);
110 let Bz_poly = self.r1cs.B.mul_vector(&Z);
111 let Cz_poly = self.r1cs.C.mul_vector(&Z);
112
113 // Prove that the
114 // Q(t) = \sum_{x \in {0, 1}^m} (Az_poly(x) * Bz_poly(x) - Cz_poly(x)) eq(t, x)
115 // is a zero-polynomial using the sum-check protocol.
116 // We evaluate Q(t) at $\tau$ and check that it is zero.
117
118 let rx = transcript.challenge_vec(m);
119
120 let sc_phase_1_timer = start_timer!(|| "Sumcheck phase 1");
121
122 let sc_phase_1 = SumCheckPhase1::new(
123 Az_poly.clone(),
124 Bz_poly.clone(),
125 Cz_poly.clone(),
126 tau.clone(),
127 rx.clone(),
128 );
129 let (sc_proof_1, (v_A, v_B, v_C)) = sc_phase_1.prove(&self.pcs_blinder, transcript);
130 end_timer!(sc_phase_1_timer);
131
132 transcript.append_fe(&v_A);
133 transcript.append_fe(&v_B);
134 transcript.append_fe(&v_C);
135
136 // Phase 2
137 let r = transcript.challenge_vec(3);
138
139 // T_2 should equal teh evaluations of the random linear combined polynomials
140
141 let ry = transcript.challenge_vec(m);
142 let sc_phase_2_timer = start_timer!(|| "Sumcheck phase 2");
143 let sc_phase_2 = SumCheckPhase2::new(
144 self.r1cs.A.clone(),
145 self.r1cs.B.clone(),
146 self.r1cs.C.clone(),
147 Z.clone(),
148 rx.clone(),
149 r.as_slice().try_into().unwrap(),
150 ry.clone(),
151 );
152
153 let sc_proof_2 = sc_phase_2.prove(&self.pcs_blinder, transcript);
154 end_timer!(sc_phase_2_timer);
155
156 let z_open_timer = start_timer!(|| "Open witness poly");
157 // Prove the evaluation of the polynomial Z(y) at ry
158 let z_eval_proof =
159 self.pcs_witness
160 .open(&committed_witness, &witness_poly, &ry[1..], transcript);
161 end_timer!(z_open_timer);
162
163 // Prove the evaluation of the polynomials A(y), B(y), C(y) at ry
164
165 let rx_ry = vec![ry, rx].concat();
166 (
167 PartialSpartanProof {
168 z_comm: witness_comm,
169 sc_proof_1,
170 sc_proof_2,
171 z_eval_proof,
172 v_A,
173 v_B,
174 v_C,
175 },
176 rx_ry,
177 )
178 }
179
180 pub fn verify_partial(
181 &self,
182 partial_proof: &PartialSpartanProof<F>,
183 transcript: &mut Transcript<F>,
184 ) {
185 partial_proof.z_comm.append_to_transcript(transcript);
186
187 let A_mle = self.r1cs.A.to_ml_extension();
188 let B_mle = self.r1cs.B.to_ml_extension();
189 let C_mle = self.r1cs.C.to_ml_extension();
190
191 let m = (self.r1cs.z_len() as f64).log2() as usize;
192 let tau = transcript.challenge_vec(m);
193 let rx = transcript.challenge_vec(m);
194
195 transcript.append_fe(&partial_proof.sc_proof_1.blinder_poly_sum);
196 transcript.append_bytes(&partial_proof.sc_proof_1.blinder_poly_eval_proof.u_hat_comm);
197
198 let rho = transcript.challenge_fe();
199
200 let ex = SumCheckPhase1::verify_round_polys(&partial_proof.sc_proof_1, &rx, rho);
201
202 self.pcs_blinder.verify(
203 &partial_proof.sc_proof_1.blinder_poly_eval_proof,
204 transcript,
205 );
206
207 // The final eval should equal
208 let v_A = partial_proof.v_A;
209 let v_B = partial_proof.v_B;
210 let v_C = partial_proof.v_C;
211
212 let T_1_eq = EqPoly::new(tau);
213 let T_1 = (v_A * v_B - v_C) * T_1_eq.eval(&rx)
214 + rho * partial_proof.sc_proof_1.blinder_poly_eval_proof.y;
215 assert_eq!(T_1, ex);
216
217 transcript.append_fe(&v_A);
218 transcript.append_fe(&v_B);
219 transcript.append_fe(&v_C);
220
221 let r = transcript.challenge_vec(3);
222 let r_A = r[0];
223 let r_B = r[1];
224 let r_C = r[2];
225
226 let ry = transcript.challenge_vec(m);
227
228 transcript.append_fe(&partial_proof.sc_proof_2.blinder_poly_sum);
229 transcript.append_bytes(&partial_proof.sc_proof_2.blinder_poly_eval_proof.u_hat_comm);
230
231 let rho_2 = transcript.challenge_fe();
232
233 let T_2 =
234 (r_A * v_A + r_B * v_B + r_C * v_C) + rho_2 * partial_proof.sc_proof_2.blinder_poly_sum;
235 let final_poly_eval =
236 SumCheckPhase2::verify_round_polys(T_2, &partial_proof.sc_proof_2, &ry);
237
238 self.pcs_blinder.verify(
239 &partial_proof.sc_proof_2.blinder_poly_eval_proof,
240 transcript,
241 );
242
243 assert_eq!(partial_proof.z_eval_proof.x, ry[1..]);
244 let rx_ry = [rx, ry.clone()].concat();
245
246 let witness_eval = partial_proof.z_eval_proof.y;
247 let A_eval = A_mle.eval(&rx_ry);
248 let B_eval = B_mle.eval(&rx_ry);
249 let C_eval = C_mle.eval(&rx_ry);
250
251 self.pcs_witness
252 .verify(&partial_proof.z_eval_proof, transcript);
253
254 let witness_len = self.r1cs.num_vars.next_power_of_two();
255 let input = (0..self.r1cs.num_input)
256 .map(|i| (witness_len + i + 1, self.r1cs.public_input[i]))
257 .collect::<Vec<(usize, F)>>();
258
259 let input_poly =
260 SparseMLPoly::new(vec![vec![(witness_len, F::ONE)], input].concat(), ry.len());
261
262 let input_poly_eval = input_poly.eval(&ry);
263
264 let z_eval = (F::ONE - ry[0]) * witness_eval + input_poly_eval;
265
266 let T_opened = (r_A * A_eval + r_B * B_eval + r_C * C_eval) * z_eval
267 + rho_2 * partial_proof.sc_proof_2.blinder_poly_eval_proof.y;
268 assert_eq!(T_opened, final_poly_eval);
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275
276 #[test]
277 fn test_shockwave_plus() {
278 type F = halo2curves::secp256k1::Fp;
279
280 let num_vars = 2usize.pow(7);
281 let num_input = 3;
282 let l = 10;
283
284 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_vars, num_input);
285
286 let num_rows = 4;
287 let ShockwavePlus = ShockwavePlus::new(r1cs.clone(), l, num_rows);
288 let mut prover_transcript = Transcript::new(b"bench");
289 let (partial_proof, _) =
290 ShockwavePlus.prove(&witness, &r1cs.public_input, &mut prover_transcript);
291
292 let mut verifier_transcript = Transcript::new(b"bench");
293 ShockwavePlus.verify_partial(&partial_proof, &mut verifier_transcript);
294 }
295}