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