1use crate::polynomial::ml_poly::MlPoly;
  2use crate::r1cs::r1cs::Matrix;
  3use crate::sumcheck::unipoly::UniPoly;
  4use crate::FieldExt;
  5use serde::{Deserialize, Serialize};
  6use tensor_pcs::{EqPoly, Transcript};
  7
  8#[derive(Serialize, Deserialize)]
  9pub struct SCPhase2Proof<F: FieldExt> {
 10    pub round_polys: Vec<UniPoly<F>>,
 11    pub blinder_poly_sum: F,
 12    pub blinder_poly_eval_claim: F,
 13}
 14
 15pub struct SumCheckPhase2<F: FieldExt> {
 16    A_mat: Matrix<F>,
 17    B_mat: Matrix<F>,
 18    C_mat: Matrix<F>,
 19    Z_evals: Vec<F>,
 20    rx: Vec<F>,
 21    r: [F; 3],
 22    challenge: Vec<F>,
 23}
 24
 25impl<F: FieldExt> SumCheckPhase2<F> {
 26    pub fn new(
 27        A_mat: Matrix<F>,
 28        B_mat: Matrix<F>,
 29        C_mat: Matrix<F>,
 30        Z_evals: Vec<F>,
 31        rx: Vec<F>,
 32        r: [F; 3],
 33        challenge: Vec<F>,
 34    ) -> Self {
 35        Self {
 36            A_mat,
 37            B_mat,
 38            C_mat,
 39            Z_evals,
 40            rx,
 41            r,
 42            challenge,
 43        }
 44    }
 45
 46    pub fn prove(&self, transcript: &mut Transcript<F>) -> SCPhase2Proof<F> {
 47        let r_A = self.r[0];
 48        let r_B = self.r[1];
 49        let r_C = self.r[2];
 50
 51        let n = self.Z_evals.len();
 52        let num_vars = (self.Z_evals.len() as f64).log2() as usize;
 53
 54        let evals_rx = EqPoly::new(self.rx.clone()).evals();
 55        let mut A_evals = vec![F::ZERO; n];
 56        let mut B_evals = vec![F::ZERO; n];
 57        let mut C_evals = vec![F::ZERO; n];
 58
 59        for entry in &self.A_mat.entries {
 60            A_evals[entry.col] += evals_rx[entry.row] * entry.val;
 61        }
 62        for entry in &self.B_mat.entries {
 63            B_evals[entry.col] += evals_rx[entry.row] * entry.val;
 64        }
 65        for entry in &self.C_mat.entries {
 66            C_evals[entry.col] += evals_rx[entry.row] * entry.val;
 67        }
 68
 69        let mut rng = rand::thread_rng();
 70        // Sample a blinding polynomial g(x_1, ..., x_m) of degree 3
 71        let random_evals = (0..2usize.pow(num_vars as u32))
 72            .map(|_| F::random(&mut rng))
 73            .collect::<Vec<F>>();
 74        let blinder_poly_sum = random_evals.iter().fold(F::ZERO, |acc, x| acc + x);
 75        let blinder_poly = MlPoly::new(random_evals);
 76
 77        transcript.append_fe(&blinder_poly_sum);
 78        let rho = transcript.challenge_fe();
 79
 80        let mut round_polys: Vec<UniPoly<F>> = Vec::<UniPoly<F>>::with_capacity(num_vars);
 81
 82        let mut A_table = A_evals.clone();
 83        let mut B_table = B_evals.clone();
 84        let mut C_table = C_evals.clone();
 85        let mut Z_table = self.Z_evals.clone();
 86        let mut blinder_table = blinder_poly.evals.clone();
 87
 88        let zero = F::ZERO;
 89        let one = F::ONE;
 90        let two = F::from(2);
 91
 92        for j in 0..num_vars {
 93            let high_index = 2usize.pow((num_vars - j - 1) as u32);
 94            let mut evals = [F::ZERO; 3];
 95
 96            for b in 0..high_index {
 97                let r_y_i = self.challenge[j];
 98                for (i, eval_at) in [zero, one, two].iter().enumerate() {
 99                    let a_eval = A_table[b] + (A_table[b + high_index] - A_table[b]) * eval_at;
100                    let b_eval = B_table[b] + (B_table[b + high_index] - B_table[b]) * eval_at;
101                    let c_eval = C_table[b] + (C_table[b + high_index] - C_table[b]) * eval_at;
102                    let z_eval = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * eval_at;
103                    let blinder_eval = blinder_table[b]
104                        + (blinder_table[b + high_index] - blinder_table[b]) * eval_at;
105                    evals[i] +=
106                        (a_eval * r_A + b_eval * r_B + c_eval * r_C) * z_eval + rho * blinder_eval;
107                }
108
109                A_table[b] = A_table[b] + (A_table[b + high_index] - A_table[b]) * r_y_i;
110                B_table[b] = B_table[b] + (B_table[b + high_index] - B_table[b]) * r_y_i;
111                C_table[b] = C_table[b] + (C_table[b + high_index] - C_table[b]) * r_y_i;
112                Z_table[b] = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * r_y_i;
113                blinder_table[b] =
114                    blinder_table[b] + (blinder_table[b + high_index] - blinder_table[b]) * r_y_i;
115            }
116
117            let round_poly = UniPoly::interpolate(&evals);
118            round_polys.push(round_poly);
119        }
120
121        let mut r_y_rev = self.challenge.clone();
122        let blinder_poly_eval_claim = blinder_poly.eval(&r_y_rev);
123
124        SCPhase2Proof {
125            round_polys,
126            blinder_poly_eval_claim,
127            blinder_poly_sum,
128        }
129    }
130
131    pub fn verify_round_polys(sum_target: F, proof: &SCPhase2Proof<F>, challenge: &[F]) -> F {
132        debug_assert_eq!(proof.round_polys.len(), challenge.len());
133
134        let zero = F::ZERO;
135        let one = F::ONE;
136
137        let mut target = sum_target;
138        for (i, round_poly) in proof.round_polys.iter().enumerate() {
139            assert_eq!(
140                round_poly.eval(zero) + round_poly.eval(one),
141                target,
142                "i = {}",
143                i
144            );
145
146            target = round_poly.eval(challenge[i]);
147        }
148
149        target
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use halo2curves::ff::Field;
157    use halo2curves::secp256k1::Fp;
158    type F = Fp;
159
160    #[test]
161    fn test_unipoly_2() {
162        let coeffs = [F::from(1u64), F::from(2u64), F::from(3u64)];
163        let eval_at = Fp::from(33);
164
165        let mut expected_eval = F::ZERO;
166        for i in 0..coeffs.len() {
167            expected_eval += coeffs[i] * eval_at.pow(&[i as u64, 0, 0, 0]);
168        }
169
170        let mut evals = [F::ZERO; 3];
171        for i in 0..3 {
172            let eval_at = F::from(i as u64);
173            let mut eval_i = F::ZERO;
174            for j in 0..coeffs.len() {
175                eval_i += coeffs[j] * eval_at.pow(&[j as u64, 0, 0, 0]);
176            }
177            evals[i] = eval_i;
178        }
179
180        let uni_poly = UniPoly::interpolate(&evals);
181        let eval = uni_poly.eval(eval_at);
182        assert_eq!(eval, expected_eval);
183    }
184}