Init commit!

shockwave_plus/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "shockwave-plus"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+ark-std = "0.4.0"
+bincode = "1.3.3"
+halo2curves = { version = "0.1.0", features = ["derive_serde"] }
+rand = "0.8.5"
+tensor-pcs = { path = "../tensor_pcs" }
+serde = { version = "1.0.152", features = ["derive"] }
+
+[dev-dependencies]
+criterion = { version = "0.4", features = ["html_reports"] }
+
+[[bench]]
+name = "prove"
+harness = false
+
+[features]
+print-trace = ["ark-std/print-trace"]

shockwave_plus/benches/prove.rs

@@ -0,0 +1,41 @@
+#![allow(non_snake_case)]
+use criterion::{criterion_group, criterion_main, Criterion};
+use shockwave_plus::ShockwavePlus;
+use shockwave_plus::R1CS;
+use tensor_pcs::Transcript;
+
+fn shockwave_plus_bench(c: &mut Criterion) {
+    type F = halo2curves::secp256k1::Fp;
+
+    for exp in [12, 15, 18] {
+        let num_cons = 2usize.pow(exp);
+        let num_vars = num_cons;
+        let num_input = 0;
+
+        let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
+
+        let mut group = c.benchmark_group(format!("ShockwavePlus num_cons: {}", num_cons));
+        let l = 319;
+        let num_rows = (((2f64 / l as f64).sqrt() * (num_vars as f64).sqrt()) as usize)
+            .next_power_of_two()
+            / 2;
+        let ShockwavePlus = ShockwavePlus::new(r1cs.clone(), l, num_rows);
+        group.bench_function("prove", |b| {
+            b.iter(|| {
+                let mut transcript = Transcript::new(b"bench");
+                ShockwavePlus.prove(&witness, &mut transcript);
+            })
+        });
+    }
+}
+
+fn set_duration() -> Criterion {
+    Criterion::default().sample_size(10)
+}
+
+criterion_group! {
+    name = benches;
+    config = set_duration();
+    targets = shockwave_plus_bench
+}
+criterion_main!(benches);

shockwave_plus/src/lib.rs

@@ -0,0 +1,272 @@
+#![allow(non_snake_case)]
+mod polynomial;
+mod r1cs;
+mod sumcheck;
+mod utils;
+
+use ark_std::{end_timer, start_timer};
+use serde::{Deserialize, Serialize};
+use sumcheck::{SCPhase1Proof, SCPhase2Proof, SumCheckPhase1, SumCheckPhase2};
+
+// Exports
+pub use r1cs::R1CS;
+pub use tensor_pcs::*;
+
+#[derive(Serialize, Deserialize)]
+pub struct PartialSpartanProof<F: FieldExt> {
+    pub z_comm: [u8; 32],
+    pub sc_proof_1: SCPhase1Proof<F>,
+    pub sc_proof_2: SCPhase2Proof<F>,
+    pub z_eval_proof: TensorMLOpening<F>,
+    pub v_A: F,
+    pub v_B: F,
+    pub v_C: F,
+}
+
+pub struct FullSpartanProof<F: FieldExt> {
+    pub partial_proof: PartialSpartanProof<F>,
+    pub A_eval_proof: TensorMLOpening<F>,
+    pub B_eval_proof: TensorMLOpening<F>,
+    pub C_eval_proof: TensorMLOpening<F>,
+}
+
+pub struct ShockwavePlus<F: FieldExt> {
+    pub r1cs: R1CS<F>,
+    pub pcs_witness: TensorMultilinearPCS<F>,
+}
+
+impl<F: FieldExt> ShockwavePlus<F> {
+    pub fn new(r1cs: R1CS<F>, l: usize, num_rows: usize) -> Self {
+        let num_cols = r1cs.num_vars / num_rows;
+
+        // Make sure that there are enough columns to run the l queries
+        assert!(num_cols > l);
+
+        let expansion_factor = 2;
+
+        let ecfft_config = rs_config::ecfft::gen_config(num_cols);
+
+        let pcs_config = TensorRSMultilinearPCSConfig::<F> {
+            expansion_factor,
+            domain_powers: None,
+            fft_domain: None,
+            ecfft_config: Some(ecfft_config),
+            l,
+            num_entries: r1cs.num_vars,
+            num_rows,
+        };
+
+        let pcs_witness = TensorMultilinearPCS::new(pcs_config);
+        Self { r1cs, pcs_witness }
+    }
+
+    pub fn prove(
+        &self,
+        witness: &[F],
+        transcript: &mut Transcript<F>,
+    ) -> (PartialSpartanProof<F>, Vec<F>) {
+        // Compute the multilinear extension of the witness
+        assert!(witness.len().is_power_of_two());
+        let witness_poly = SparseMLPoly::from_dense(witness.to_vec());
+
+        // Commit the witness polynomial
+        let comm_witness_timer = start_timer!(|| "Commit witness");
+        let committed_witness = self.pcs_witness.commit(&witness_poly);
+        let witness_comm = committed_witness.committed_tree.root;
+        end_timer!(comm_witness_timer);
+
+        transcript.append_bytes(&witness_comm);
+
+        // ############################
+        // Phase 1: The sum-checks
+        // ###################
+
+        let m = (self.r1cs.num_vars as f64).log2() as usize;
+        let tau = transcript.challenge_vec(m);
+        let mut tau_rev = tau.clone();
+        tau_rev.reverse();
+
+        // First
+        // Compute the multilinear extension of the R1CS matrices.
+        // Prove that he Q_poly is a zero-polynomial
+
+        // Q_poly is a zero-polynomial iff F_io evaluates to zero
+        // over the m-dimensional boolean hypercube..
+
+        // We prove using the sum-check protocol.
+
+        // G_poly = A_poly * B_poly - C_poly
+
+        let num_rows = self.r1cs.num_cons;
+        let Az_poly = self.r1cs.A.mul_vector(num_rows, witness);
+        let Bz_poly = self.r1cs.B.mul_vector(num_rows, witness);
+        let Cz_poly = self.r1cs.C.mul_vector(num_rows, witness);
+
+        // Prove that the polynomial Q(t)
+        // \sum_{x \in {0, 1}^m} (Az_poly(x) * Bz_poly(x) - Cz_poly(x)) eq(tau, x)
+        // is a zero-polynomial using the sum-check protocol.
+
+        let rx = transcript.challenge_vec(m);
+        let mut rx_rev = rx.clone();
+        rx_rev.reverse();
+
+        let sc_phase_1_timer = start_timer!(|| "Sumcheck phase 1");
+
+        let sc_phase_1 = SumCheckPhase1::new(
+            Az_poly.clone(),
+            Bz_poly.clone(),
+            Cz_poly.clone(),
+            tau_rev.clone(),
+            rx.clone(),
+        );
+        let (sc_proof_1, (v_A, v_B, v_C)) = sc_phase_1.prove(transcript);
+        end_timer!(sc_phase_1_timer);
+
+        transcript.append_fe(&v_A);
+        transcript.append_fe(&v_B);
+        transcript.append_fe(&v_C);
+
+        // Phase 2
+        let r = transcript.challenge_vec(3);
+
+        // T_2 should equal teh evaluations of the random linear combined polynomials
+
+        let ry = transcript.challenge_vec(m);
+        let sc_phase_2_timer = start_timer!(|| "Sumcheck phase 2");
+        let sc_phase_2 = SumCheckPhase2::new(
+            self.r1cs.A.clone(),
+            self.r1cs.B.clone(),
+            self.r1cs.C.clone(),
+            witness.to_vec(),
+            rx.clone(),
+            r.as_slice().try_into().unwrap(),
+            ry.clone(),
+        );
+
+        let sc_proof_2 = sc_phase_2.prove(transcript);
+        end_timer!(sc_phase_2_timer);
+
+        let mut ry_rev = ry.clone();
+        ry_rev.reverse();
+        let z_open_timer = start_timer!(|| "Open witness poly");
+        // Prove the evaluation of the polynomial Z(y) at ry
+        let z_eval_proof =
+            self.pcs_witness
+                .open(&committed_witness, &witness_poly, &ry_rev, transcript);
+        end_timer!(z_open_timer);
+
+        // Prove the evaluation of the polynomials A(y), B(y), C(y) at ry
+
+        let rx_ry = vec![ry_rev, rx_rev].concat();
+        (
+            PartialSpartanProof {
+                z_comm: witness_comm,
+                sc_proof_1,
+                sc_proof_2,
+                z_eval_proof,
+                v_A,
+                v_B,
+                v_C,
+            },
+            rx_ry,
+        )
+    }
+
+    pub fn verify_partial(
+        &self,
+        partial_proof: &PartialSpartanProof<F>,
+        transcript: &mut Transcript<F>,
+    ) {
+        partial_proof.z_comm.append_to_transcript(transcript);
+
+        let A_mle = self.r1cs.A.to_ml_extension();
+        let B_mle = self.r1cs.B.to_ml_extension();
+        let C_mle = self.r1cs.C.to_ml_extension();
+
+        let m = (self.r1cs.num_vars as f64).log2() as usize;
+        let tau = transcript.challenge_vec(m);
+        let rx = transcript.challenge_vec(m);
+        let mut rx_rev = rx.clone();
+        rx_rev.reverse();
+
+        transcript.append_fe(&partial_proof.sc_proof_1.blinder_poly_sum);
+        let rho = transcript.challenge_fe();
+
+        let ex = SumCheckPhase1::verify_round_polys(&partial_proof.sc_proof_1, &rx, rho);
+
+        // The final eval should equal
+        let v_A = partial_proof.v_A;
+        let v_B = partial_proof.v_B;
+        let v_C = partial_proof.v_C;
+
+        let T_1_eq = EqPoly::new(tau);
+        let T_1 = (v_A * v_B - v_C) * T_1_eq.eval(&rx_rev)
+            + rho * partial_proof.sc_proof_1.blinder_poly_eval_claim;
+        assert_eq!(T_1, ex);
+
+        transcript.append_fe(&v_A);
+        transcript.append_fe(&v_B);
+        transcript.append_fe(&v_C);
+
+        let r = transcript.challenge_vec(3);
+        let r_A = r[0];
+        let r_B = r[1];
+        let r_C = r[2];
+
+        let ry = transcript.challenge_vec(m);
+
+        transcript.append_fe(&partial_proof.sc_proof_2.blinder_poly_sum);
+        let rho_2 = transcript.challenge_fe();
+
+        let T_2 =
+            (r_A * v_A + r_B * v_B + r_C * v_C) + rho_2 * partial_proof.sc_proof_2.blinder_poly_sum;
+        let final_poly_eval =
+            SumCheckPhase2::verify_round_polys(T_2, &partial_proof.sc_proof_2, &ry);
+
+        let mut ry_rev = ry.clone();
+        ry_rev.reverse();
+
+        let rx_ry = [rx, ry].concat();
+        assert_eq!(partial_proof.z_eval_proof.x, ry_rev);
+
+        let z_eval = partial_proof.z_eval_proof.y;
+        let A_eval = A_mle.eval(&rx_ry);
+        let B_eval = B_mle.eval(&rx_ry);
+        let C_eval = C_mle.eval(&rx_ry);
+
+        self.pcs_witness.verify(
+            &partial_proof.z_eval_proof,
+            &partial_proof.z_comm,
+            transcript,
+        );
+
+        let T_opened = (r_A * A_eval + r_B * B_eval + r_C * C_eval) * z_eval
+            + rho_2 * partial_proof.sc_proof_2.blinder_poly_eval_claim;
+        assert_eq!(T_opened, final_poly_eval);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_shockwave_plus() {
+        type F = halo2curves::secp256k1::Fp;
+
+        let num_cons = 2usize.pow(6);
+        let num_vars = num_cons;
+        let num_input = 0;
+        let l = 10;
+
+        let (r1cs, witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
+
+        let num_rows = 4;
+        let ShockwavePlus = ShockwavePlus::new(r1cs.clone(), l, num_rows);
+        let mut prover_transcript = Transcript::new(b"bench");
+        let (partial_proof, _) = ShockwavePlus.prove(&witness, &mut prover_transcript);
+
+        let mut verifier_transcript = Transcript::new(b"bench");
+        ShockwavePlus.verify_partial(&partial_proof, &mut verifier_transcript);
+    }
+}

shockwave_plus/src/polynomial/blinder_poly.rs

@@ -0,0 +1,34 @@
+use crate::FieldExt;
+
+pub struct BlinderPoly<F: FieldExt> {
+    inner_poly_coeffs: Vec<Vec<F>>,
+}
+
+impl<F: FieldExt> BlinderPoly<F> {
+    pub fn sample_random(num_vars: usize, degree: usize) -> Self {
+        let mut rng = rand::thread_rng();
+        let inner_poly_coeffs = (0..num_vars)
+            .map(|_| (0..(degree + 1)).map(|_| F::random(&mut rng)).collect())
+            .collect();
+
+        Self { inner_poly_coeffs }
+    }
+
+    pub fn eval(&self, x: &[F]) -> F {
+        let mut res = F::ZERO;
+
+        for (coeffs, x_i) in self.inner_poly_coeffs.iter().zip(x.iter()) {
+            let mut tmp = F::ZERO;
+            let mut x_i_pow = F::ONE;
+
+            for coeff in coeffs.iter() {
+                tmp += *coeff * x_i_pow;
+                x_i_pow *= x_i;
+            }
+
+            res += tmp;
+        }
+
+        res
+    }
+}

shockwave_plus/src/polynomial/ml_poly.rs

@@ -0,0 +1,39 @@
+use tensor_pcs::EqPoly;
+
+use crate::FieldExt;
+
+#[derive(Clone, Debug)]
+pub struct MlPoly<F> {
+    pub evals: Vec<F>,
+    pub num_vars: usize,
+}
+
+impl<F: FieldExt> MlPoly<F> {
+    pub fn new(evals: Vec<F>) -> Self {
+        assert!(evals.len().is_power_of_two());
+        let num_vars = (evals.len() as f64).log2() as usize;
+        Self { evals, num_vars }
+    }
+
+    fn dot_prod(x: &[F], y: &[F]) -> F {
+        assert_eq!(x.len(), y.len());
+        let mut result = F::ZERO;
+        for i in 0..x.len() {
+            result += x[i] * y[i];
+        }
+        result
+    }
+
+    // Evaluate the multilinear extension of the polynomial `a`, at point `t`.
+    // `a` is in evaluation form.
+    pub fn eval(&self, t: &[F]) -> F {
+        let n = self.evals.len();
+        debug_assert_eq!((n as f64).log2() as usize, t.len());
+        // Evaluate the multilinear extension of the polynomial `a`,
+        // over the boolean hypercube
+
+        let eq_evals = EqPoly::new(t.to_vec()).evals();
+
+        Self::dot_prod(&self.evals, &eq_evals)
+    }
+}

shockwave_plus/src/polynomial/mod.rs

@@ -0,0 +1,2 @@
+pub mod blinder_poly;
+pub mod ml_poly;

shockwave_plus/src/r1cs/mod.rs

@@ -0,0 +1,3 @@
+pub mod r1cs;
+
+pub use r1cs::R1CS;

shockwave_plus/src/r1cs/r1cs.rs

@@ -0,0 +1,306 @@
+use crate::FieldExt;
+use halo2curves::ff::Field;
+use tensor_pcs::SparseMLPoly;
+
+#[derive(Clone)]
+pub struct SparseMatrixEntry<F: FieldExt> {
+    pub row: usize,
+    pub col: usize,
+    pub val: F,
+}
+
+#[derive(Clone)]
+pub struct Matrix<F: FieldExt> {
+    pub entries: Vec<SparseMatrixEntry<F>>,
+    pub num_cols: usize,
+    pub num_rows: usize,
+}
+
+impl<F> Matrix<F>
+where
+    F: FieldExt,
+{
+    pub fn new(entries: Vec<SparseMatrixEntry<F>>, num_cols: usize, num_rows: usize) -> Self {
+        assert!((num_cols * num_rows).is_power_of_two());
+        Self {
+            entries,
+            num_cols,
+            num_rows,
+        }
+    }
+
+    pub fn mul_vector(&self, num_rows: usize, vec: &[F]) -> Vec<F> {
+        let mut result = vec![F::ZERO; num_rows];
+        let entries = &self.entries;
+        for i in 0..entries.len() {
+            let row = entries[i].row;
+            let col = entries[i].col;
+            let val = entries[i].val;
+            result[row] += val * vec[col];
+        }
+        result
+    }
+
+    // Return a multilinear extension of the matrix
+    // with num_vars * num_vars entries
+    pub fn to_ml_extension(&self) -> SparseMLPoly<F> {
+        let mut evals = Vec::with_capacity(self.entries.len());
+        let entries = &self.entries;
+        let num_cols = self.num_cols;
+        for i in 0..entries.len() {
+            let row = entries[i].row;
+            let col = entries[i].col;
+            let val = entries[i].val;
+            evals.push(((row * num_cols) + col, val));
+        }
+        let ml_poly_num_vars = ((self.num_cols * self.num_rows) as f64).log2() as usize;
+        let ml_poly = SparseMLPoly::new(evals, ml_poly_num_vars);
+        ml_poly
+    }
+
+    /*
+    pub fn fast_to_coeffs(&self, s: usize, x: F) -> Vec<F> {
+        let mut result = F::ZERO;
+        for entry in &self.0 {
+            let row = entry.0;
+            let col = entry.1;
+            let val = entry.2;
+
+            let index = row * 2usize.pow(s as u32) + col;
+            // Get the degrees of the nonzero coefficients
+            // Tensor product (1 - x_0)(1 - x_1)
+            let base = index;
+            let zero_bits = degree & !base;
+
+            let mut zero_bit_degrees = vec![];
+            for j in 0..s {
+                if zero_bits & (1 << j) != 0 {
+                    zero_bit_degrees.push(j);
+                }
+            }
+
+            let mut term = val;
+            for degree in zero_bit_degrees {
+                term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
+            }
+            result += term;
+        }
+        result
+    }
+     */
+
+    /*
+    pub fn fast_uni_eval(&self, s: usize, x: F) -> F {
+        let degree = 2usize.pow(s as u32);
+
+        let mut result = F::ZERO;
+        for entry in &self.0 {
+            let row = entry.0;
+            let col = entry.1;
+            let val = entry.2;
+
+            let index = row * 2usize.pow(s as u32) + col;
+            // Get the degrees of the nonzero coefficients
+            // Tensor product (1 - x_0)(1 - x_1)
+            let base = index;
+            let zero_bits = degree & !base;
+
+            let mut zero_bit_degrees = vec![];
+            for j in 0..s {
+                if zero_bits & (1 << j) != 0 {
+                    zero_bit_degrees.push(j);
+                }
+            }
+
+            let mut term = val;
+            for degree in zero_bit_degrees {
+                term *= x.pow(&[base as u64, 0, 0, 0]) - x.pow(&[(degree + base) as u64, 0, 0, 0]);
+            }
+            result += term;
+        }
+        result
+    }
+     */
+}
+
+#[derive(Clone)]
+pub struct R1CS<F>
+where
+    F: FieldExt,
+{
+    pub A: Matrix<F>,
+    pub B: Matrix<F>,
+    pub C: Matrix<F>,
+    pub public_input: Vec<F>,
+    pub num_cons: usize,
+    pub num_vars: usize,
+    pub num_input: usize,
+}
+
+impl<F> R1CS<F>
+where
+    F: FieldExt,
+{
+    pub fn hadamard_prod(a: &[F], b: &[F]) -> Vec<F> {
+        assert_eq!(a.len(), b.len());
+        let mut result = vec![F::ZERO; a.len()];
+        for i in 0..a.len() {
+            result[i] = a[i] * b[i];
+        }
+        result
+    }
+
+    pub fn produce_synthetic_r1cs(
+        num_cons: usize,
+        num_vars: usize,
+        num_input: usize,
+    ) -> (Self, Vec<F>) {
+        //        assert_eq!(num_cons, num_vars);
+        let mut public_input = Vec::with_capacity(num_input);
+        let mut witness = Vec::with_capacity(num_vars);
+
+        for i in 0..num_input {
+            public_input.push(F::from((i + 1) as u64));
+        }
+
+        for i in 0..num_vars {
+            witness.push(F::from((i + 1) as u64));
+        }
+
+        let z: Vec<F> = vec![public_input.clone(), witness.clone()].concat();
+
+        let mut A_entries: Vec<SparseMatrixEntry<F>> = vec![];
+        let mut B_entries: Vec<SparseMatrixEntry<F>> = vec![];
+        let mut C_entries: Vec<SparseMatrixEntry<F>> = vec![];
+
+        for i in 0..num_cons {
+            let A_col = i % num_vars;
+            let B_col = (i + 1) % num_vars;
+            let C_col = (i + 2) % num_vars;
+
+            // For the i'th constraint,
+            // add the value 1 at the (i % num_vars)th column of A, B.
+            // Compute the corresponding C_column value so that A_i * B_i = C_i
+            // we apply multiplication since the Hadamard product is computed for Az ・ Bz,
+
+            // We only _enable_ a single variable in each constraint.
+            A_entries.push(SparseMatrixEntry {
+                row: i,
+                col: A_col,
+                val: F::ONE,
+            });
+            B_entries.push(SparseMatrixEntry {
+                row: i,
+                col: B_col,
+                val: F::ONE,
+            });
+            C_entries.push(SparseMatrixEntry {
+                row: i,
+                col: C_col,
+                val: (z[A_col] * z[B_col]) * z[C_col].invert().unwrap(),
+            });
+        }
+
+        let A = Matrix::new(A_entries, num_vars, num_cons);
+
+        let B = Matrix::new(B_entries, num_vars, num_cons);
+
+        let C = Matrix::new(C_entries, num_vars, num_cons);
+
+        (
+            Self {
+                A,
+                B,
+                C,
+                public_input,
+                num_cons,
+                num_vars,
+                num_input,
+            },
+            witness,
+        )
+    }
+
+    pub fn is_sat(&self, witness: &Vec<F>, public_input: &Vec<F>) -> bool {
+        let mut z = Vec::with_capacity(witness.len() + public_input.len() + 1);
+        z.extend(public_input);
+        z.extend(witness);
+
+        let Az = self.A.mul_vector(self.num_cons, &z);
+        let Bz = self.B.mul_vector(self.num_cons, &z);
+        let Cz = self.C.mul_vector(self.num_cons, &z);
+
+        Self::hadamard_prod(&Az, &Bz) == Cz
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::utils::boolean_hypercube;
+
+    use super::*;
+    type F = halo2curves::secp256k1::Fp;
+    use crate::polynomial::ml_poly::MlPoly;
+
+    #[test]
+    fn test_r1cs() {
+        let num_cons = 2usize.pow(5);
+        let num_vars = num_cons;
+        let num_input = 0;
+
+        let (r1cs, mut witness) = R1CS::<F>::produce_synthetic_r1cs(num_cons, num_vars, num_input);
+
+        assert_eq!(witness.len(), num_vars);
+        assert_eq!(r1cs.public_input.len(), num_input);
+
+        assert!(r1cs.is_sat(&witness, &r1cs.public_input));
+
+        // Should assert if the witness is invalid
+        witness[0] = witness[0] + F::one();
+        assert!(r1cs.is_sat(&r1cs.public_input, &witness) == false);
+        witness[0] = witness[0] - F::one();
+
+        /*
+        // Should assert if the public input is invalid
+        let mut public_input = r1cs.public_input.clone();
+        public_input[0] = public_input[0] + F::one();
+        assert!(r1cs.is_sat(&witness, &public_input) == false);
+         */
+
+        // Test MLE
+        let s = (num_vars as f64).log2() as usize;
+        let A_mle = r1cs.A.to_ml_extension();
+        let B_mle = r1cs.B.to_ml_extension();
+        let C_mle = r1cs.C.to_ml_extension();
+        let Z_mle = MlPoly::new(witness);
+
+        for c in &boolean_hypercube(s) {
+            let mut eval_a = F::zero();
+            let mut eval_b = F::zero();
+            let mut eval_c = F::zero();
+            for b in &boolean_hypercube(s) {
+                let mut b_rev = b.clone();
+                b_rev.reverse();
+                let z_eval = Z_mle.eval(&b_rev);
+                let mut eval_matrix = [b.as_slice(), c.as_slice()].concat();
+                eval_matrix.reverse();
+                eval_a += A_mle.eval(&eval_matrix) * z_eval;
+                eval_b += B_mle.eval(&eval_matrix) * z_eval;
+                eval_c += C_mle.eval(&eval_matrix) * z_eval;
+            }
+            let eval_con = eval_a * eval_b - eval_c;
+            assert_eq!(eval_con, F::zero());
+        }
+    }
+
+    /*
+    #[test]
+    fn test_fast_uni_eval() {
+        let (r1cs, _) = R1CS::<F>::produce_synthetic_r1cs(8, 8, 0);
+
+        let eval_at = F::from(33);
+        let result = r1cs.A.fast_uni_eval(r1cs.num_vars, eval_at);
+        println!("result: {:?}", result);
+    }
+     */
+}

shockwave_plus/src/sumcheck/mod.rs

@@ -0,0 +1,6 @@
+mod sc_phase_1;
+mod sc_phase_2;
+pub mod unipoly;
+
+pub use sc_phase_1::{SCPhase1Proof, SumCheckPhase1};
+pub use sc_phase_2::{SCPhase2Proof, SumCheckPhase2};

shockwave_plus/src/sumcheck/sc_phase_1.rs

@@ -0,0 +1,175 @@
+use crate::polynomial::ml_poly::MlPoly;
+use crate::sumcheck::unipoly::UniPoly;
+use serde::{Deserialize, Serialize};
+use tensor_pcs::{EqPoly, Transcript};
+
+use crate::FieldExt;
+
+#[derive(Serialize, Deserialize)]
+pub struct SCPhase1Proof<F: FieldExt> {
+    pub blinder_poly_sum: F,
+    pub blinder_poly_eval_claim: F,
+    pub round_polys: Vec<UniPoly<F>>,
+}
+
+pub struct SumCheckPhase1<F: FieldExt> {
+    Az_evals: Vec<F>,
+    Bz_evals: Vec<F>,
+    Cz_evals: Vec<F>,
+    bound_eq_poly: EqPoly<F>,
+    challenge: Vec<F>,
+}
+
+impl<F: FieldExt> SumCheckPhase1<F> {
+    pub fn new(
+        Az_evals: Vec<F>,
+        Bz_evals: Vec<F>,
+        Cz_evals: Vec<F>,
+        tau: Vec<F>,
+        challenge: Vec<F>,
+    ) -> Self {
+        let bound_eq_poly = EqPoly::new(tau);
+        Self {
+            Az_evals,
+            Bz_evals,
+            Cz_evals,
+            bound_eq_poly,
+            challenge,
+        }
+    }
+
+    pub fn prove(&self, transcript: &mut Transcript<F>) -> (SCPhase1Proof<F>, (F, F, F)) {
+        let num_vars = (self.Az_evals.len() as f64).log2() as usize;
+        let mut round_polys = Vec::<UniPoly<F>>::with_capacity(num_vars - 1);
+
+        let mut rng = rand::thread_rng();
+        // Sample a blinding polynomial g(x_1, ..., x_m) of degree 3
+        let random_evals = (0..2usize.pow(num_vars as u32))
+            .map(|_| F::random(&mut rng))
+            .collect::<Vec<F>>();
+        let blinder_poly_sum = random_evals.iter().fold(F::ZERO, |acc, x| acc + x);
+        let blinder_poly = MlPoly::new(random_evals);
+
+        transcript.append_fe(&blinder_poly_sum);
+        let rho = transcript.challenge_fe();
+
+        // Compute the sum of g(x_1, ... x_m) over the boolean hypercube
+
+        // Do the sum check for f + \rho g
+
+        let mut A_table = self.Az_evals.clone();
+        let mut B_table = self.Bz_evals.clone();
+        let mut C_table = self.Cz_evals.clone();
+        let mut blinder_table = blinder_poly.evals.clone();
+        let mut eq_table = self.bound_eq_poly.evals();
+
+        let zero = F::ZERO;
+        let one = F::ONE;
+        let two = F::from(2);
+        let three = F::from(3);
+
+        for j in 0..num_vars {
+            let r_i = self.challenge[j];
+
+            let high_index = 2usize.pow((num_vars - j - 1) as u32);
+
+            let mut evals = [F::ZERO; 4];
+
+            // https://eprint.iacr.org/2019/317.pdf#subsection.3.2
+            for b in 0..high_index {
+                for (i, eval_at) in [zero, one, two, three].iter().enumerate() {
+                    let a_eval = A_table[b] + (A_table[b + high_index] - A_table[b]) * eval_at;
+                    let b_eval = B_table[b] + (B_table[b + high_index] - B_table[b]) * eval_at;
+                    let c_eval = C_table[b] + (C_table[b + high_index] - C_table[b]) * eval_at;
+                    let eq_eval = eq_table[b] + (eq_table[b + high_index] - eq_table[b]) * eval_at;
+                    let blinder_eval = blinder_table[b]
+                        + (blinder_table[b + high_index] - blinder_table[b]) * eval_at;
+                    evals[i] += ((a_eval * b_eval - c_eval) * eq_eval) + rho * blinder_eval;
+                }
+
+                A_table[b] = A_table[b] + (A_table[b + high_index] - A_table[b]) * r_i;
+                B_table[b] = B_table[b] + (B_table[b + high_index] - B_table[b]) * r_i;
+                C_table[b] = C_table[b] + (C_table[b + high_index] - C_table[b]) * r_i;
+                eq_table[b] = eq_table[b] + (eq_table[b + high_index] - eq_table[b]) * r_i;
+                blinder_table[b] =
+                    blinder_table[b] + (blinder_table[b + high_index] - blinder_table[b]) * r_i;
+            }
+
+            let round_poly = UniPoly::interpolate(&evals);
+
+            round_polys.push(round_poly);
+        }
+
+        let v_A = A_table[0];
+        let v_B = B_table[0];
+        let v_C = C_table[0];
+
+        let rx = self.challenge.clone();
+        let blinder_poly_eval_claim = blinder_poly.eval(&rx);
+
+        // Prove the evaluation of the blinder polynomial at rx.
+
+        (
+            SCPhase1Proof {
+                blinder_poly_sum,
+                round_polys,
+                blinder_poly_eval_claim,
+            },
+            (v_A, v_B, v_C),
+        )
+    }
+
+    pub fn verify_round_polys(proof: &SCPhase1Proof<F>, challenge: &[F], rho: F) -> F {
+        debug_assert_eq!(proof.round_polys.len(), challenge.len());
+
+        let zero = F::ZERO;
+        let one = F::ONE;
+
+        // target = 0 + rho * blinder_poly_sum
+        let mut target = rho * proof.blinder_poly_sum;
+        for (i, round_poly) in proof.round_polys.iter().enumerate() {
+            assert_eq!(
+                round_poly.eval(zero) + round_poly.eval(one),
+                target,
+                "round poly {} failed",
+                i
+            );
+            target = round_poly.eval(challenge[i]);
+        }
+
+        target
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use halo2curves::secp256k1::Fp;
+    type F = Fp;
+    use halo2curves::ff::Field;
+
+    #[test]
+    fn test_unipoly_3() {
+        let coeffs = [F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)];
+        let eval_at = Fp::from(33);
+
+        let mut expected_eval = F::ZERO;
+        for i in 0..coeffs.len() {
+            expected_eval += coeffs[i] * eval_at.pow(&[3 - i as u64, 0, 0, 0]);
+        }
+
+        let mut evals = [F::ZERO; 4];
+        for i in 0..4 {
+            let eval_at = F::from(i as u64);
+            let mut eval_i = F::ZERO;
+            for j in 0..coeffs.len() {
+                eval_i += coeffs[j] * eval_at.pow(&[3 - j as u64, 0, 0, 0]);
+            }
+            evals[i] = eval_i;
+        }
+
+        let uni_poly = UniPoly::interpolate(&evals);
+        let eval = uni_poly.eval(eval_at);
+        assert_eq!(eval, expected_eval);
+    }
+}

shockwave_plus/src/sumcheck/sc_phase_2.rs

@@ -0,0 +1,184 @@
+use crate::polynomial::ml_poly::MlPoly;
+use crate::r1cs::r1cs::Matrix;
+use crate::sumcheck::unipoly::UniPoly;
+use crate::FieldExt;
+use serde::{Deserialize, Serialize};
+use tensor_pcs::{EqPoly, Transcript};
+
+#[derive(Serialize, Deserialize)]
+pub struct SCPhase2Proof<F: FieldExt> {
+    pub round_polys: Vec<UniPoly<F>>,
+    pub blinder_poly_sum: F,
+    pub blinder_poly_eval_claim: F,
+}
+
+pub struct SumCheckPhase2<F: FieldExt> {
+    A_mat: Matrix<F>,
+    B_mat: Matrix<F>,
+    C_mat: Matrix<F>,
+    Z_evals: Vec<F>,
+    rx: Vec<F>,
+    r: [F; 3],
+    challenge: Vec<F>,
+}
+
+impl<F: FieldExt> SumCheckPhase2<F> {
+    pub fn new(
+        A_mat: Matrix<F>,
+        B_mat: Matrix<F>,
+        C_mat: Matrix<F>,
+        Z_evals: Vec<F>,
+        rx: Vec<F>,
+        r: [F; 3],
+        challenge: Vec<F>,
+    ) -> Self {
+        Self {
+            A_mat,
+            B_mat,
+            C_mat,
+            Z_evals,
+            rx,
+            r,
+            challenge,
+        }
+    }
+
+    pub fn prove(&self, transcript: &mut Transcript<F>) -> SCPhase2Proof<F> {
+        let r_A = self.r[0];
+        let r_B = self.r[1];
+        let r_C = self.r[2];
+
+        let n = self.Z_evals.len();
+        let num_vars = (self.Z_evals.len() as f64).log2() as usize;
+
+        let evals_rx = EqPoly::new(self.rx.clone()).evals();
+        let mut A_evals = vec![F::ZERO; n];
+        let mut B_evals = vec![F::ZERO; n];
+        let mut C_evals = vec![F::ZERO; n];
+
+        for entry in &self.A_mat.entries {
+            A_evals[entry.col] += evals_rx[entry.row] * entry.val;
+        }
+        for entry in &self.B_mat.entries {
+            B_evals[entry.col] += evals_rx[entry.row] * entry.val;
+        }
+        for entry in &self.C_mat.entries {
+            C_evals[entry.col] += evals_rx[entry.row] * entry.val;
+        }
+
+        let mut rng = rand::thread_rng();
+        // Sample a blinding polynomial g(x_1, ..., x_m) of degree 3
+        let random_evals = (0..2usize.pow(num_vars as u32))
+            .map(|_| F::random(&mut rng))
+            .collect::<Vec<F>>();
+        let blinder_poly_sum = random_evals.iter().fold(F::ZERO, |acc, x| acc + x);
+        let blinder_poly = MlPoly::new(random_evals);
+
+        transcript.append_fe(&blinder_poly_sum);
+        let rho = transcript.challenge_fe();
+
+        let mut round_polys: Vec<UniPoly<F>> = Vec::<UniPoly<F>>::with_capacity(num_vars);
+
+        let mut A_table = A_evals.clone();
+        let mut B_table = B_evals.clone();
+        let mut C_table = C_evals.clone();
+        let mut Z_table = self.Z_evals.clone();
+        let mut blinder_table = blinder_poly.evals.clone();
+
+        let zero = F::ZERO;
+        let one = F::ONE;
+        let two = F::from(2);
+
+        for j in 0..num_vars {
+            let high_index = 2usize.pow((num_vars - j - 1) as u32);
+            let mut evals = [F::ZERO; 3];
+
+            for b in 0..high_index {
+                let r_y_i = self.challenge[j];
+                for (i, eval_at) in [zero, one, two].iter().enumerate() {
+                    let a_eval = A_table[b] + (A_table[b + high_index] - A_table[b]) * eval_at;
+                    let b_eval = B_table[b] + (B_table[b + high_index] - B_table[b]) * eval_at;
+                    let c_eval = C_table[b] + (C_table[b + high_index] - C_table[b]) * eval_at;
+                    let z_eval = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * eval_at;
+                    let blinder_eval = blinder_table[b]
+                        + (blinder_table[b + high_index] - blinder_table[b]) * eval_at;
+                    evals[i] +=
+                        (a_eval * r_A + b_eval * r_B + c_eval * r_C) * z_eval + rho * blinder_eval;
+                }
+
+                A_table[b] = A_table[b] + (A_table[b + high_index] - A_table[b]) * r_y_i;
+                B_table[b] = B_table[b] + (B_table[b + high_index] - B_table[b]) * r_y_i;
+                C_table[b] = C_table[b] + (C_table[b + high_index] - C_table[b]) * r_y_i;
+                Z_table[b] = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * r_y_i;
+                blinder_table[b] =
+                    blinder_table[b] + (blinder_table[b + high_index] - blinder_table[b]) * r_y_i;
+            }
+
+            let round_poly = UniPoly::interpolate(&evals);
+            round_polys.push(round_poly);
+        }
+
+        let mut r_y_rev = self.challenge.clone();
+        let blinder_poly_eval_claim = blinder_poly.eval(&r_y_rev);
+
+        SCPhase2Proof {
+            round_polys,
+            blinder_poly_eval_claim,
+            blinder_poly_sum,
+        }
+    }
+
+    pub fn verify_round_polys(sum_target: F, proof: &SCPhase2Proof<F>, challenge: &[F]) -> F {
+        debug_assert_eq!(proof.round_polys.len(), challenge.len());
+
+        let zero = F::ZERO;
+        let one = F::ONE;
+
+        let mut target = sum_target;
+        for (i, round_poly) in proof.round_polys.iter().enumerate() {
+            assert_eq!(
+                round_poly.eval(zero) + round_poly.eval(one),
+                target,
+                "i = {}",
+                i
+            );
+
+            target = round_poly.eval(challenge[i]);
+        }
+
+        target
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use halo2curves::ff::Field;
+    use halo2curves::secp256k1::Fp;
+    type F = Fp;
+
+    #[test]
+    fn test_unipoly_2() {
+        let coeffs = [F::from(1u64), F::from(2u64), F::from(3u64)];
+        let eval_at = Fp::from(33);
+
+        let mut expected_eval = F::ZERO;
+        for i in 0..coeffs.len() {
+            expected_eval += coeffs[i] * eval_at.pow(&[i as u64, 0, 0, 0]);
+        }
+
+        let mut evals = [F::ZERO; 3];
+        for i in 0..3 {
+            let eval_at = F::from(i as u64);
+            let mut eval_i = F::ZERO;
+            for j in 0..coeffs.len() {
+                eval_i += coeffs[j] * eval_at.pow(&[j as u64, 0, 0, 0]);
+            }
+            evals[i] = eval_i;
+        }
+
+        let uni_poly = UniPoly::interpolate(&evals);
+        let eval = uni_poly.eval(eval_at);
+        assert_eq!(eval, expected_eval);
+    }
+}

shockwave_plus/src/sumcheck/unipoly.rs

@@ -0,0 +1,81 @@
+use crate::FieldExt;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+pub struct UniPoly<F: FieldExt> {
+    pub coeffs: Vec<F>,
+}
+
+impl<F: FieldExt> UniPoly<F> {
+    fn eval_cubic(&self, x: F) -> F {
+        // ax^3 + bx^2 + cx + d
+        let x_sq = x.square();
+        let x_cub = x_sq * x;
+
+        let a = self.coeffs[0];
+        let b = self.coeffs[1];
+        let c = self.coeffs[2];
+        let d = self.coeffs[3];
+
+        a * x_cub + b * x_sq + c * x + d
+    }
+
+    fn eval_quadratic(&self, x: F) -> F {
+        // ax^3 + bx^2 + cx + d
+        let x_sq = x.square();
+
+        let a = self.coeffs[0];
+        let b = self.coeffs[1];
+        let c = self.coeffs[2];
+
+        a * x_sq + b * x + c
+    }
+
+    pub fn eval(&self, x: F) -> F {
+        if self.coeffs.len() == 3 {
+            self.eval_quadratic(x)
+        } else {
+            self.eval_cubic(x)
+        }
+    }
+
+    pub fn interpolate(evals: &[F]) -> Self {
+        debug_assert!(
+            evals.len() == 4 || evals.len() == 3,
+            "Only cubic and quadratic polynomials are supported"
+        );
+
+        let two_inv = F::TWO_INV;
+
+        if evals.len() == 4 {
+            // ax^3 + bx^2 + cx + d
+            let six_inv = F::from(6u64).invert().unwrap();
+
+            let d = evals[0];
+            let a = six_inv
+                * (evals[3] - evals[2] - evals[2] - evals[2] + evals[1] + evals[1] + evals[1]
+                    - evals[0]);
+            let b = two_inv
+                * (evals[0] + evals[0] - evals[1] - evals[1] - evals[1] - evals[1] - evals[1]
+                    + evals[2]
+                    + evals[2]
+                    + evals[2]
+                    + evals[2]
+                    - evals[3]);
+
+            let c = evals[1] - d - a - b;
+
+            Self {
+                coeffs: vec![a, b, c, d],
+            }
+        } else {
+            let c = evals[0];
+            let a = (evals[2] - evals[1] - evals[1] + evals[0]) * two_inv;
+            let b = evals[1] - a - c;
+
+            Self {
+                coeffs: vec![a, b, c],
+            }
+        }
+    }
+}

shockwave_plus/src/utils.rs

@@ -0,0 +1,19 @@
+use crate::FieldExt;
+
+// Returns a vector of vectors of length m, where each vector is a boolean vector (little endian)
+pub fn boolean_hypercube<F: FieldExt>(m: usize) -> Vec<Vec<F>> {
+    let n = 2usize.pow(m as u32);
+
+    let mut boolean_hypercube = Vec::<Vec<F>>::with_capacity(n);
+
+    for i in 0..n {
+        let mut tmp = Vec::with_capacity(m);
+        for j in 0..m {
+            let i_b = F::from((i >> j & 1) as u64);
+            tmp.push(i_b);
+        }
+        boolean_hypercube.push(tmp);
+    }
+
+    boolean_hypercube
+}