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