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