1use crate::r1cs::r1cs::Matrix;
  2use crate::sumcheck::unipoly::UniPoly;
  3use crate::FieldExt;
  4use serde::{Deserialize, Serialize};
  5use tensor_pcs::{EqPoly, SparseMLPoly, 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 random_evals = (0..2usize.pow(num_vars as u32))
 75            .map(|_| F::random(&mut rng))
 76            .collect::<Vec<F>>();
 77        let blinder_poly_sum = random_evals.iter().fold(F::ZERO, |acc, x| acc + x);
 78        let blinder_poly = SparseMLPoly::from_dense(random_evals);
 79        let blinder_poly_comm = pcs.commit(&blinder_poly);
 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
 93            .evals
 94            .iter()
 95            .map(|(_, x)| *x)
 96            .collect::<Vec<F>>();
 97
 98        let zero = F::ZERO;
 99        let one = F::ONE;
100        let two = F::from(2);
101
102        for j in 0..num_vars {
103            let high_index = 2usize.pow((num_vars - j - 1) as u32);
104            let mut evals = [F::ZERO; 3];
105
106            for b in 0..high_index {
107                let r_y_i = self.challenge[j];
108                for (i, eval_at) in [zero, one, two].iter().enumerate() {
109                    let a_eval = A_table[b] + (A_table[b + high_index] - A_table[b]) * eval_at;
110                    let b_eval = B_table[b] + (B_table[b + high_index] - B_table[b]) * eval_at;
111                    let c_eval = C_table[b] + (C_table[b + high_index] - C_table[b]) * eval_at;
112                    let z_eval = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * eval_at;
113                    let blinder_eval = blinder_table[b]
114                        + (blinder_table[b + high_index] - blinder_table[b]) * eval_at;
115                    evals[i] +=
116                        (a_eval * r_A + b_eval * r_B + c_eval * r_C) * z_eval + rho * blinder_eval;
117                }
118
119                A_table[b] = A_table[b] + (A_table[b + high_index] - A_table[b]) * r_y_i;
120                B_table[b] = B_table[b] + (B_table[b + high_index] - B_table[b]) * r_y_i;
121                C_table[b] = C_table[b] + (C_table[b + high_index] - C_table[b]) * r_y_i;
122                Z_table[b] = Z_table[b] + (Z_table[b + high_index] - Z_table[b]) * r_y_i;
123                blinder_table[b] =
124                    blinder_table[b] + (blinder_table[b + high_index] - blinder_table[b]) * r_y_i;
125            }
126
127            let round_poly = UniPoly::interpolate(&evals);
128            round_polys.push(round_poly);
129        }
130
131        let ry = self.challenge.clone();
132
133        let blinder_poly_eval_proof = pcs.open(&blinder_poly_comm, &blinder_poly, &ry, transcript);
134
135        SCPhase2Proof {
136            round_polys,
137            blinder_poly_sum,
138            blinder_poly_eval_proof,
139        }
140    }
141
142    pub fn verify_round_polys(sum_target: F, proof: &SCPhase2Proof<F>, challenge: &[F]) -> F {
143        debug_assert_eq!(proof.round_polys.len(), challenge.len());
144
145        let zero = F::ZERO;
146        let one = F::ONE;
147
148        let mut target = sum_target;
149        for (i, round_poly) in proof.round_polys.iter().enumerate() {
150            assert_eq!(
151                round_poly.eval(zero) + round_poly.eval(one),
152                target,
153                "i = {}",
154                i
155            );
156
157            target = round_poly.eval(challenge[i]);
158        }
159
160        target
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use halo2curves::ff::Field;
168    use halo2curves::secp256k1::Fp;
169    type F = Fp;
170
171    #[test]
172    fn test_unipoly_2() {
173        let coeffs = [F::from(1u64), F::from(2u64), F::from(3u64)];
174        let eval_at = Fp::from(33);
175
176        let mut expected_eval = F::ZERO;
177        for i in 0..coeffs.len() {
178            expected_eval += coeffs[i] * eval_at.pow(&[i as u64, 0, 0, 0]);
179        }
180
181        let mut evals = [F::ZERO; 3];
182        for i in 0..3 {
183            let eval_at = F::from(i as u64);
184            let mut eval_i = F::ZERO;
185            for j in 0..coeffs.len() {
186                eval_i += coeffs[j] * eval_at.pow(&[j as u64, 0, 0, 0]);
187            }
188            evals[i] = eval_i;
189        }
190
191        let uni_poly = UniPoly::interpolate(&evals);
192        let eval = uni_poly.eval(eval_at);
193        assert_eq!(eval, expected_eval);
194    }
195}