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