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