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