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