1use crate::utils::espresso::sum_check::SumCheck;
  2use crate::utils::virtual_polynomial::VPAuxInfo;
  3use crate::{
  4    transcript::{poseidon::PoseidonTranscript, TranscriptVar},
  5    utils::sum_check::{structs::IOPProof, IOPSumCheck},
  6};
  7use ark_crypto_primitives::sponge::Absorb;
  8use ark_ec::{CurveGroup, Group};
  9/// Heavily inspired from testudo: https://github.com/cryptonetlab/testudo/tree/master
 10/// Some changes:
 11/// - Typings to better stick to ark_poly's API
 12/// - Uses `folding-schemes`' own `TranscriptVar` trait and `PoseidonTranscriptVar` struct
 13/// - API made closer to gadgets found in `folding-schemes`
 14use ark_ff::PrimeField;
 15use ark_poly::{univariate::DensePolynomial, DenseUVPolynomial};
 16use ark_r1cs_std::{
 17    alloc::{AllocVar, AllocationMode},
 18    eq::EqGadget,
 19    fields::fp::FpVar,
 20};
 21use ark_relations::r1cs::{Namespace, SynthesisError};
 22use std::{borrow::Borrow, marker::PhantomData};
 23
 24#[derive(Clone, Debug)]
 25pub struct DensePolynomialVar<F: PrimeField> {
 26    pub coeffs: Vec<FpVar<F>>,
 27}
 28
 29impl<F: PrimeField> AllocVar<DensePolynomial<F>, F> for DensePolynomialVar<F> {
 30    fn new_variable<T: Borrow<DensePolynomial<F>>>(
 31        cs: impl Into<Namespace<F>>,
 32        f: impl FnOnce() -> Result<T, SynthesisError>,
 33        mode: AllocationMode,
 34    ) -> Result<Self, SynthesisError> {
 35        f().and_then(|c| {
 36            let cs = cs.into();
 37            let cp: &DensePolynomial<F> = c.borrow();
 38            let mut coeffs_var = Vec::<FpVar<F>>::with_capacity(cp.coeffs.len());
 39            for coeff in cp.coeffs.iter() {
 40                let coeff_var = FpVar::<F>::new_variable(cs.clone(), || Ok(coeff), mode)?;
 41                coeffs_var.push(coeff_var);
 42            }
 43            Ok(Self { coeffs: coeffs_var })
 44        })
 45    }
 46}
 47
 48impl<F: PrimeField> DensePolynomialVar<F> {
 49    pub fn eval_at_zero(&self) -> FpVar<F> {
 50        self.coeffs[0].clone()
 51    }
 52
 53    pub fn eval_at_one(&self) -> FpVar<F> {
 54        let mut res = self.coeffs[0].clone();
 55        for i in 1..self.coeffs.len() {
 56            res = &res + &self.coeffs[i];
 57        }
 58        res
 59    }
 60
 61    pub fn evaluate(&self, r: &FpVar<F>) -> FpVar<F> {
 62        let mut eval = self.coeffs[0].clone();
 63        let mut power = r.clone();
 64
 65        for i in 1..self.coeffs.len() {
 66            eval += &power * &self.coeffs[i];
 67            power *= r;
 68        }
 69        eval
 70    }
 71}
 72
 73#[derive(Clone, Debug)]
 74pub struct IOPProofVar<C: CurveGroup> {
 75    // We have to be generic over a CurveGroup because instantiating a IOPProofVar will call IOPSumCheck which requires a CurveGroup
 76    pub proofs: Vec<DensePolynomialVar<C::ScalarField>>,
 77    pub claim: FpVar<C::ScalarField>,
 78}
 79
 80impl<C: CurveGroup> AllocVar<IOPProof<C::ScalarField>, C::ScalarField> for IOPProofVar<C>
 81where
 82    <C as Group>::ScalarField: Absorb,
 83{
 84    fn new_variable<T: Borrow<IOPProof<C::ScalarField>>>(
 85        cs: impl Into<Namespace<C::ScalarField>>,
 86        f: impl FnOnce() -> Result<T, SynthesisError>,
 87        mode: AllocationMode,
 88    ) -> Result<Self, SynthesisError> {
 89        f().and_then(|c| {
 90            let cs = cs.into();
 91            let cp: &IOPProof<C::ScalarField> = c.borrow();
 92            let claim = IOPSumCheck::<C, PoseidonTranscript<C>>::extract_sum(cp);
 93            let claim = FpVar::<C::ScalarField>::new_variable(cs.clone(), || Ok(claim), mode)?;
 94            let mut proofs =
 95                Vec::<DensePolynomialVar<C::ScalarField>>::with_capacity(cp.proofs.len());
 96            for proof in cp.proofs.iter() {
 97                let poly = DensePolynomial::from_coefficients_slice(&proof.coeffs);
 98                let proof = DensePolynomialVar::<C::ScalarField>::new_variable(
 99                    cs.clone(),
100                    || Ok(poly),
101                    mode,
102                )?;
103                proofs.push(proof);
104            }
105            Ok(Self { proofs, claim })
106        })
107    }
108}
109
110#[derive(Clone, Debug)]
111pub struct VPAuxInfoVar<F: PrimeField> {
112    pub num_variables: FpVar<F>,
113    pub max_degree: FpVar<F>,
114}
115
116impl<F: PrimeField> AllocVar<VPAuxInfo<F>, F> for VPAuxInfoVar<F> {
117    fn new_variable<T: Borrow<VPAuxInfo<F>>>(
118        cs: impl Into<Namespace<F>>,
119        f: impl FnOnce() -> Result<T, SynthesisError>,
120        mode: AllocationMode,
121    ) -> Result<Self, SynthesisError> {
122        f().and_then(|c| {
123            let cs = cs.into();
124            let cp: &VPAuxInfo<F> = c.borrow();
125            let num_variables = FpVar::<F>::new_variable(
126                cs.clone(),
127                || Ok(F::from(cp.num_variables as u64)),
128                mode,
129            )?;
130            let max_degree =
131                FpVar::<F>::new_variable(cs.clone(), || Ok(F::from(cp.max_degree as u64)), mode)?;
132            Ok(Self {
133                num_variables,
134                max_degree,
135            })
136        })
137    }
138}
139
140#[derive(Debug, Clone)]
141pub struct SumCheckVerifierGadget<C: CurveGroup> {
142    _f: PhantomData<C>,
143}
144
145impl<C: CurveGroup> SumCheckVerifierGadget<C> {
146    #[allow(clippy::type_complexity)]
147    pub fn verify(
148        iop_proof_var: &IOPProofVar<C>,
149        poly_aux_info_var: &VPAuxInfoVar<C::ScalarField>,
150        transcript_var: &mut impl TranscriptVar<C::ScalarField>,
151    ) -> Result<(Vec<FpVar<C::ScalarField>>, Vec<FpVar<C::ScalarField>>), SynthesisError> {
152        let mut e_vars = vec![iop_proof_var.claim.clone()];
153        let mut r_vars: Vec<FpVar<C::ScalarField>> = Vec::new();
154        transcript_var.absorb(poly_aux_info_var.num_variables.clone())?;
155        transcript_var.absorb(poly_aux_info_var.max_degree.clone())?;
156
157        for poly_var in iop_proof_var.proofs.iter() {
158            let res = poly_var.eval_at_one() + poly_var.eval_at_zero();
159            let e_var = e_vars.last().ok_or(SynthesisError::Unsatisfiable)?;
160            res.enforce_equal(e_var)?;
161            transcript_var.absorb_vec(&poly_var.coeffs)?;
162            let r_i_var = transcript_var.get_challenge()?;
163            e_vars.push(poly_var.evaluate(&r_i_var));
164            r_vars.push(r_i_var);
165        }
166
167        Ok((e_vars, r_vars))
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use crate::{
174        folding::circuits::sum_check::{IOPProofVar, VPAuxInfoVar},
175        transcript::{
176            poseidon::{poseidon_test_config, PoseidonTranscript, PoseidonTranscriptVar},
177            Transcript, TranscriptVar,
178        },
179        utils::{
180            sum_check::{structs::IOPProof, IOPSumCheck, SumCheck},
181            virtual_polynomial::VirtualPolynomial,
182        },
183    };
184    use ark_crypto_primitives::sponge::{poseidon::PoseidonConfig, Absorb};
185    use ark_ec::CurveGroup;
186    use ark_ff::Field;
187    use ark_pallas::{Fr, Projective};
188    use ark_poly::{
189        univariate::DensePolynomial, DenseMultilinearExtension, DenseUVPolynomial,
190        MultilinearExtension, Polynomial,
191    };
192    use ark_r1cs_std::{alloc::AllocVar, R1CSVar};
193    use ark_relations::r1cs::ConstraintSystem;
194    use std::sync::Arc;
195
196    use super::SumCheckVerifierGadget;
197
198    pub type TestSumCheckProof<F> = (VirtualPolynomial<F>, PoseidonConfig<F>, IOPProof<F>);
199
200    /// Primarily used for testing the sumcheck gadget
201    /// Returns a random virtual polynomial, the poseidon config used and the associated sumcheck proof
202    pub fn get_test_sumcheck_proof<C: CurveGroup>(
203        num_vars: usize,
204    ) -> TestSumCheckProof<C::ScalarField>
205    where
206        <C as ark_ec::Group>::ScalarField: Absorb,
207    {
208        let mut rng = ark_std::test_rng();
209        let poseidon_config: PoseidonConfig<C::ScalarField> =
210            poseidon_test_config::<C::ScalarField>();
211        let mut poseidon_transcript_prove = PoseidonTranscript::<C>::new(&poseidon_config);
212        let poly_mle = DenseMultilinearExtension::rand(num_vars, &mut rng);
213        let virtual_poly =
214            VirtualPolynomial::new_from_mle(&Arc::new(poly_mle), C::ScalarField::ONE);
215        let sum_check: IOPProof<C::ScalarField> = IOPSumCheck::<C, PoseidonTranscript<C>>::prove(
216            &virtual_poly,
217            &mut poseidon_transcript_prove,
218        )
219        .unwrap();
220        (virtual_poly, poseidon_config, sum_check)
221    }
222
223    #[test]
224    fn test_sum_check_circuit() {
225        for num_vars in 1..15 {
226            let cs = ConstraintSystem::<Fr>::new_ref();
227            let (virtual_poly, poseidon_config, sum_check) =
228                get_test_sumcheck_proof::<Projective>(num_vars);
229            let mut poseidon_var: PoseidonTranscriptVar<Fr> =
230                PoseidonTranscriptVar::new(cs.clone(), &poseidon_config);
231            let iop_proof_var =
232                IOPProofVar::<Projective>::new_witness(cs.clone(), || Ok(&sum_check)).unwrap();
233            let poly_aux_info_var =
234                VPAuxInfoVar::<Fr>::new_witness(cs.clone(), || Ok(virtual_poly.aux_info)).unwrap();
235            let res = SumCheckVerifierGadget::<Projective>::verify(
236                &iop_proof_var,
237                &poly_aux_info_var,
238                &mut poseidon_var,
239            );
240
241            assert!(res.is_ok());
242            let (circuit_evals, r_challenges) = res.unwrap();
243
244            // 1. assert claim from circuit is equal to the one from the sum-check
245            let claim: Fr =
246                IOPSumCheck::<Projective, PoseidonTranscript<Projective>>::extract_sum(&sum_check);
247            assert_eq!(circuit_evals[0].value().unwrap(), claim);
248
249            // 2. assert that all in-circuit evaluations are equal to the ones from the sum-check
250            for ((proof, point), circuit_eval) in sum_check
251                .proofs
252                .iter()
253                .zip(sum_check.point.iter())
254                .zip(circuit_evals.iter().skip(1))
255            // we skip the first one since it's the above checked claim
256            {
257                let poly = DensePolynomial::from_coefficients_slice(&proof.coeffs);
258                let eval = poly.evaluate(point);
259                assert_eq!(eval, circuit_eval.value().unwrap());
260            }
261
262            // 3. assert that all challenges are equal to the ones from the sum-check
263            for (point, r_challenge) in sum_check.point.iter().zip(r_challenges.iter()) {
264                assert_eq!(*point, r_challenge.value().unwrap());
265            }
266
267            assert!(cs.is_satisfied().unwrap());
268        }
269    }
270}