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}
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);
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 Self { r1cs, pcs_witness }
61 }
62
63 pub fn prove(
64 &self,
65 r1cs_witness: &[F],
66 transcript: &mut Transcript<F>,
67 ) -> (PartialSpartanProof<F>, Vec<F>) {
68 // Compute the multilinear extension of the witness
69 assert!(r1cs_witness.len().is_power_of_two());
70 let witness_poly = SparseMLPoly::from_dense(r1cs_witness.to_vec());
71
72 // Commit the witness polynomial
73 let comm_witness_timer = start_timer!(|| "Commit witness");
74 let committed_witness = self.pcs_witness.commit(&witness_poly);
75 let witness_comm = committed_witness.committed_tree.root;
76 end_timer!(comm_witness_timer);
77
78 // Add the witness commitment to the transcript
79 transcript.append_bytes(&witness_comm);
80
81 // ############################
82 // Phase 1
83 // ###################
84
85 let m = (self.r1cs.num_vars as f64).log2() as usize;
86 let tau = transcript.challenge_vec(m);
87 let mut tau_rev = tau.clone();
88 tau_rev.reverse();
89
90 let num_rows = self.r1cs.num_cons;
91 let Az_poly = self.r1cs.A.mul_vector(num_rows, r1cs_witness);
92 let Bz_poly = self.r1cs.B.mul_vector(num_rows, r1cs_witness);
93 let Cz_poly = self.r1cs.C.mul_vector(num_rows, r1cs_witness);
94
95 // Prove that the
96 // Q(t) = \sum_{x \in {0, 1}^m} (Az_poly(x) * Bz_poly(x) - Cz_poly(x)) eq(t, x)
97 // is a zero-polynomial using the sum-check protocol.
98 // We evaluate Q(t) at $\tau$ and check that it is zero.
99
100 let rx = transcript.challenge_vec(m);
101 let mut rx_rev = rx.clone();
102 rx_rev.reverse();
103
104 let sc_phase_1_timer = start_timer!(|| "Sumcheck phase 1");
105
106 let sc_phase_1 = SumCheckPhase1::new(
107 Az_poly.clone(),
108 Bz_poly.clone(),
109 Cz_poly.clone(),
110 tau_rev.clone(),
111 rx.clone(),
112 );
113 let (sc_proof_1, (v_A, v_B, v_C)) = sc_phase_1.prove(&self.pcs_witness, transcript);
114 end_timer!(sc_phase_1_timer);
115
116 transcript.append_fe(&v_A);
117 transcript.append_fe(&v_B);
118 transcript.append_fe(&v_C);
119
120 // Phase 2
121 let r = transcript.challenge_vec(3);
122
123 // T_2 should equal teh evaluations of the random linear combined polynomials
124
125 let ry = transcript.challenge_vec(m);
126 let sc_phase_2_timer = start_timer!(|| "Sumcheck phase 2");
127 let sc_phase_2 = SumCheckPhase2::new(
128 self.r1cs.A.clone(),
129 self.r1cs.B.clone(),
130 self.r1cs.C.clone(),
131 r1cs_witness.to_vec(),
132 rx.clone(),
133 r.as_slice().try_into().unwrap(),
134 ry.clone(),
135 );
136
137 let sc_proof_2 = sc_phase_2.prove(&self.pcs_witness, transcript);
138 end_timer!(sc_phase_2_timer);
139
140 let mut ry_rev = ry.clone();
141 ry_rev.reverse();
142 let z_open_timer = start_timer!(|| "Open witness poly");
143 // Prove the evaluation of the polynomial Z(y) at ry
144 let z_eval_proof =
145 self.pcs_witness
146 .open(&committed_witness, &witness_poly, &ry_rev, transcript);
147 end_timer!(z_open_timer);
148
149 // Prove the evaluation of the polynomials A(y), B(y), C(y) at ry
150
151 let rx_ry = vec![ry_rev, rx_rev].concat();
152 (
153 PartialSpartanProof {
154 z_comm: witness_comm,
155 sc_proof_1,
156 sc_proof_2,
157 z_eval_proof,
158 v_A,
159 v_B,
160 v_C,
161 },
162 rx_ry,
163 )
164 }
165
166 pub fn verify_partial(
167 &self,
168 partial_proof: &PartialSpartanProof<F>,
169 transcript: &mut Transcript<F>,
170 ) {
171 partial_proof.z_comm.append_to_transcript(transcript);
172
173 let A_mle = self.r1cs.A.to_ml_extension();
174 let B_mle = self.r1cs.B.to_ml_extension();
175 let C_mle = self.r1cs.C.to_ml_extension();
176
177 let m = (self.r1cs.num_vars as f64).log2() as usize;
178 let tau = transcript.challenge_vec(m);
179 let rx = transcript.challenge_vec(m);
180 let mut rx_rev = rx.clone();
181 rx_rev.reverse();
182
183 transcript.append_fe(&partial_proof.sc_proof_1.blinder_poly_sum);
184 transcript.append_bytes(&partial_proof.sc_proof_1.blinder_poly_eval_proof.u_hat_comm);
185
186 let rho = transcript.challenge_fe();
187
188 let ex = SumCheckPhase1::verify_round_polys(&partial_proof.sc_proof_1, &rx, rho);
189
190 self.pcs_witness.verify(
191 &partial_proof.sc_proof_1.blinder_poly_eval_proof,
192 transcript,
193 );
194
195 // The final eval should equal
196 let v_A = partial_proof.v_A;
197 let v_B = partial_proof.v_B;
198 let v_C = partial_proof.v_C;
199
200 let T_1_eq = EqPoly::new(tau);
201 let T_1 = (v_A * v_B - v_C) * T_1_eq.eval(&rx_rev)
202 + rho * partial_proof.sc_proof_1.blinder_poly_eval_proof.y;
203 assert_eq!(T_1, ex);
204
205 transcript.append_fe(&v_A);
206 transcript.append_fe(&v_B);
207 transcript.append_fe(&v_C);
208
209 let r = transcript.challenge_vec(3);
210 let r_A = r[0];
211 let r_B = r[1];
212 let r_C = r[2];
213
214 let ry = transcript.challenge_vec(m);
215
216 transcript.append_fe(&partial_proof.sc_proof_2.blinder_poly_sum);
217 transcript.append_bytes(&partial_proof.sc_proof_2.blinder_poly_eval_proof.u_hat_comm);
218
219 let rho_2 = transcript.challenge_fe();
220
221 let T_2 =
222 (r_A * v_A + r_B * v_B + r_C * v_C) + rho_2 * partial_proof.sc_proof_2.blinder_poly_sum;
223 let final_poly_eval =
224 SumCheckPhase2::verify_round_polys(T_2, &partial_proof.sc_proof_2, &ry);
225
226 self.pcs_witness.verify(
227 &partial_proof.sc_proof_2.blinder_poly_eval_proof,
228 transcript,
229 );
230
231 let mut ry_rev = ry.clone();
232 ry_rev.reverse();
233
234 let rx_ry = [rx, ry].concat();
235 assert_eq!(partial_proof.z_eval_proof.x, ry_rev);
236
237 let z_eval = partial_proof.z_eval_proof.y;
238 let A_eval = A_mle.eval(&rx_ry);
239 let B_eval = B_mle.eval(&rx_ry);
240 let C_eval = C_mle.eval(&rx_ry);
241
242 self.pcs_witness
243 .verify(&partial_proof.z_eval_proof, transcript);
244
245 let T_opened = (r_A * A_eval + r_B * B_eval + r_C * C_eval) * z_eval
246 + rho_2 * partial_proof.sc_proof_2.blinder_poly_eval_proof.y;
247 assert_eq!(T_opened, final_poly_eval);
248 }
249}
250
251#[cfg(test)]
252mod tests {
253 use super::*;
254
255 #[test]
256 fn test_shockwave_plus() {
257 type F = halo2curves::secp256k1::Fp;
258
259 let num_cons = 2usize.pow(6);
260 let num_vars = num_cons;
261 let num_input = 0;
262 let l = 10;
263
264 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
265
266 let num_rows = 4;
267 let ShockwavePlus = ShockwavePlus::new(r1cs.clone(), l, num_rows);
268 let mut prover_transcript = Transcript::new(b"bench");
269 let (partial_proof, _) = ShockwavePlus.prove(&witness, &mut prover_transcript);
270
271 let mut verifier_transcript = Transcript::new(b"bench");
272 ShockwavePlus.verify_partial(&partial_proof, &mut verifier_transcript);
273 }
274}