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 witness: &[F],
66 transcript: &mut Transcript<F>,
67 ) -> (PartialSpartanProof<F>, Vec<F>) {
68 // Compute the multilinear extension of the witness
69 assert!(witness.len().is_power_of_two());
70 let witness_poly = SparseMLPoly::from_dense(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 transcript.append_bytes(&witness_comm);
79
80 // ############################
81 // Phase 1: The sum-checks
82 // ###################
83
84 let m = (self.r1cs.num_vars as f64).log2() as usize;
85 let tau = transcript.challenge_vec(m);
86 let mut tau_rev = tau.clone();
87 tau_rev.reverse();
88
89 // First
90 // Compute the multilinear extension of the R1CS matrices.
91 // Prove that he Q_poly is a zero-polynomial
92
93 // Q_poly is a zero-polynomial iff F_io evaluates to zero
94 // over the m-dimensional boolean hypercube..
95
96 // We prove using the sum-check protocol.
97
98 // G_poly = A_poly * B_poly - C_poly
99
100 let num_rows = self.r1cs.num_cons;
101 let Az_poly = self.r1cs.A.mul_vector(num_rows, witness);
102 let Bz_poly = self.r1cs.B.mul_vector(num_rows, witness);
103 let Cz_poly = self.r1cs.C.mul_vector(num_rows, witness);
104
105 // Prove that the polynomial Q(t)
106 // \sum_{x \in {0, 1}^m} (Az_poly(x) * Bz_poly(x) - Cz_poly(x)) eq(tau, x)
107 // is a zero-polynomial using the sum-check protocol.
108
109 let rx = transcript.challenge_vec(m);
110 let mut rx_rev = rx.clone();
111 rx_rev.reverse();
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_rev.clone(),
120 rx.clone(),
121 );
122 let (sc_proof_1, (v_A, v_B, v_C)) = sc_phase_1.prove(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 witness.to_vec(),
141 rx.clone(),
142 r.as_slice().try_into().unwrap(),
143 ry.clone(),
144 );
145
146 let sc_proof_2 = sc_phase_2.prove(transcript);
147 end_timer!(sc_phase_2_timer);
148
149 let mut ry_rev = ry.clone();
150 ry_rev.reverse();
151 let z_open_timer = start_timer!(|| "Open witness poly");
152 // Prove the evaluation of the polynomial Z(y) at ry
153 let z_eval_proof =
154 self.pcs_witness
155 .open(&committed_witness, &witness_poly, &ry_rev, transcript);
156 end_timer!(z_open_timer);
157
158 // Prove the evaluation of the polynomials A(y), B(y), C(y) at ry
159
160 let rx_ry = vec![ry_rev, rx_rev].concat();
161 (
162 PartialSpartanProof {
163 z_comm: witness_comm,
164 sc_proof_1,
165 sc_proof_2,
166 z_eval_proof,
167 v_A,
168 v_B,
169 v_C,
170 },
171 rx_ry,
172 )
173 }
174
175 pub fn verify_partial(
176 &self,
177 partial_proof: &PartialSpartanProof<F>,
178 transcript: &mut Transcript<F>,
179 ) {
180 partial_proof.z_comm.append_to_transcript(transcript);
181
182 let A_mle = self.r1cs.A.to_ml_extension();
183 let B_mle = self.r1cs.B.to_ml_extension();
184 let C_mle = self.r1cs.C.to_ml_extension();
185
186 let m = (self.r1cs.num_vars as f64).log2() as usize;
187 let tau = transcript.challenge_vec(m);
188 let rx = transcript.challenge_vec(m);
189 let mut rx_rev = rx.clone();
190 rx_rev.reverse();
191
192 transcript.append_fe(&partial_proof.sc_proof_1.blinder_poly_sum);
193 let rho = transcript.challenge_fe();
194
195 let ex = SumCheckPhase1::verify_round_polys(&partial_proof.sc_proof_1, &rx, rho);
196
197 // The final eval should equal
198 let v_A = partial_proof.v_A;
199 let v_B = partial_proof.v_B;
200 let v_C = partial_proof.v_C;
201
202 let T_1_eq = EqPoly::new(tau);
203 let T_1 = (v_A * v_B - v_C) * T_1_eq.eval(&rx_rev)
204 + rho * partial_proof.sc_proof_1.blinder_poly_eval_claim;
205 assert_eq!(T_1, ex);
206
207 transcript.append_fe(&v_A);
208 transcript.append_fe(&v_B);
209 transcript.append_fe(&v_C);
210
211 let r = transcript.challenge_vec(3);
212 let r_A = r[0];
213 let r_B = r[1];
214 let r_C = r[2];
215
216 let ry = transcript.challenge_vec(m);
217
218 transcript.append_fe(&partial_proof.sc_proof_2.blinder_poly_sum);
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 let mut ry_rev = ry.clone();
227 ry_rev.reverse();
228
229 let rx_ry = [rx, ry].concat();
230 assert_eq!(partial_proof.z_eval_proof.x, ry_rev);
231
232 let z_eval = partial_proof.z_eval_proof.y;
233 let A_eval = A_mle.eval(&rx_ry);
234 let B_eval = B_mle.eval(&rx_ry);
235 let C_eval = C_mle.eval(&rx_ry);
236
237 self.pcs_witness.verify(
238 &partial_proof.z_eval_proof,
239 &partial_proof.z_comm,
240 transcript,
241 );
242
243 let T_opened = (r_A * A_eval + r_B * B_eval + r_C * C_eval) * z_eval
244 + rho_2 * partial_proof.sc_proof_2.blinder_poly_eval_claim;
245 assert_eq!(T_opened, final_poly_eval);
246 }
247}
248
249#[cfg(test)]
250mod tests {
251 use super::*;
252
253 #[test]
254 fn test_shockwave_plus() {
255 type F = halo2curves::secp256k1::Fp;
256
257 let num_cons = 2usize.pow(6);
258 let num_vars = num_cons;
259 let num_input = 0;
260 let l = 10;
261
262 let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
263
264 let num_rows = 4;
265 let ShockwavePlus = ShockwavePlus::new(r1cs.clone(), l, num_rows);
266 let mut prover_transcript = Transcript::new(b"bench");
267 let (partial_proof, _) = ShockwavePlus.prove(&witness, &mut prover_transcript);
268
269 let mut verifier_transcript = Transcript::new(b"bench");
270 ShockwavePlus.verify_partial(&partial_proof, &mut verifier_transcript);
271 }
272}