1//! This module implements the sum check protocol.
  2//! Currently this is a simple wrapper of the sumcheck protocol
  3//! from Arkworks.
  4
  5use crate::{
  6    errors::PolyIOPErrors,
  7    structs::{DomainInfo, IOPProof, IOPProverState, IOPVerifierState, SubClaim},
  8    transcript::IOPTranscript,
  9    virtual_poly::VirtualPolynomial,
 10    PolyIOP,
 11};
 12use ark_ff::PrimeField;
 13use ark_std::{end_timer, start_timer};
 14
 15mod prover;
 16mod verifier;
 17
 18/// Trait for doing sum check protocols.
 19pub trait SumCheck<F: PrimeField> {
 20    type Proof;
 21    type PolyList;
 22    type DomainInfo;
 23    type SubClaim;
 24    type Transcript;
 25
 26    /// extract sum from the proof
 27    fn extract_sum(proof: &Self::Proof) -> F;
 28
 29    /// Initialize the system with a transcript
 30    ///
 31    /// This function is optional -- in the case where a SumCheck is
 32    /// an building block for a more complex protocol, the transcript
 33    /// may be initialized by this complex protocol, and passed to the
 34    /// SumCheck prover/verifier.
 35    fn init_transcript() -> Self::Transcript;
 36
 37    /// Generate proof of the sum of polynomial over {0,1}^`num_vars`
 38    ///
 39    /// The polynomial is represented in the form of a VirtualPolynomial.
 40    fn prove(
 41        poly: &Self::PolyList,
 42        transcript: &mut Self::Transcript,
 43    ) -> Result<Self::Proof, PolyIOPErrors>;
 44
 45    /// Verify the claimed sum using the proof
 46    fn verify(
 47        sum: F,
 48        proof: &Self::Proof,
 49        domain_info: &Self::DomainInfo,
 50        transcript: &mut Self::Transcript,
 51    ) -> Result<Self::SubClaim, PolyIOPErrors>;
 52}
 53
 54/// Trait for sum check protocol prover side APIs.
 55pub trait SumCheckProver<F: PrimeField>
 56where
 57    Self: Sized,
 58{
 59    type PolyList;
 60    type ProverMessage;
 61
 62    /// Initialize the prover to argue for the sum of polynomial over
 63    /// {0,1}^`num_vars`
 64    ///
 65    /// The polynomial is represented in the form of a VirtualPolynomial.
 66    fn prover_init(polynomial: &Self::PolyList) -> Result<Self, PolyIOPErrors>;
 67
 68    /// receive message from verifier, generate prover message, and proceed to
 69    /// next round
 70    ///
 71    /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
 72    fn prove_round_and_update_state(
 73        &mut self,
 74        challenge: &Option<F>,
 75    ) -> Result<Self::ProverMessage, PolyIOPErrors>;
 76}
 77
 78/// Trait for sum check protocol verifier side APIs.
 79pub trait SumCheckVerifier<F: PrimeField> {
 80    type DomainInfo;
 81    type ProverMessage;
 82    type Challenge;
 83    type Transcript;
 84    type SubClaim;
 85
 86    /// initialize the verifier
 87    fn verifier_init(index_info: &Self::DomainInfo) -> Self;
 88
 89    /// Run verifier at current round, given prover message
 90    ///
 91    /// Normally, this function should perform actual verification. Instead,
 92    /// `verify_round` only samples and stores randomness and perform
 93    /// verifications altogether in `check_and_generate_subclaim` at
 94    /// the last step.
 95    fn verify_round_and_update_state(
 96        &mut self,
 97        prover_msg: &Self::ProverMessage,
 98        transcript: &mut Self::Transcript,
 99    ) -> Result<Self::Challenge, PolyIOPErrors>;
100
101    /// verify the sumcheck phase, and generate the subclaim
102    ///
103    /// If the asserted sum is correct, then the multilinear polynomial
104    /// evaluated at `subclaim.point` is `subclaim.expected_evaluation`.
105    /// Otherwise, it is highly unlikely that those two will be equal.
106    /// Larger field size guarantees smaller soundness error.
107    fn check_and_generate_subclaim(
108        &self,
109        asserted_sum: &F,
110    ) -> Result<Self::SubClaim, PolyIOPErrors>;
111}
112
113impl<F: PrimeField> SumCheck<F> for PolyIOP<F> {
114    type Proof = IOPProof<F>;
115
116    type PolyList = VirtualPolynomial<F>;
117
118    type DomainInfo = DomainInfo<F>;
119
120    type SubClaim = SubClaim<F>;
121
122    type Transcript = IOPTranscript<F>;
123
124    fn extract_sum(proof: &Self::Proof) -> F {
125        let start = start_timer!(|| "extract sum");
126        let res = proof.proofs[0].evaluations[0] + proof.proofs[0].evaluations[1];
127        end_timer!(start);
128        res
129    }
130
131    /// Initialize the system with a transcript
132    ///
133    /// This function is optional -- in the case where a SumCheck is
134    /// an building block for a more complex protocol, the transcript
135    /// may be initialized by this complex protocol, and passed to the
136    /// SumCheck prover/verifier.
137    fn init_transcript() -> Self::Transcript {
138        let start = start_timer!(|| "init transcript");
139        let res = IOPTranscript::<F>::new(b"Initializing SumCheck transcript");
140        end_timer!(start);
141        res
142    }
143
144    /// Generate proof of the sum of polynomial over {0,1}^`num_vars`
145    ///
146    /// The polynomial is represented in the form of a VirtualPolynomial.
147    fn prove(
148        poly: &Self::PolyList,
149        transcript: &mut Self::Transcript,
150    ) -> Result<Self::Proof, PolyIOPErrors> {
151        let start = start_timer!(|| "sum check prove");
152
153        transcript.append_domain_info(&poly.domain_info)?;
154
155        let mut prover_state = IOPProverState::prover_init(poly)?;
156        let mut challenge = None;
157        let mut prover_msgs = Vec::with_capacity(poly.domain_info.num_variables);
158        for _ in 0..poly.domain_info.num_variables {
159            let prover_msg =
160                IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge)?;
161            transcript.append_prover_message(&prover_msg)?;
162            prover_msgs.push(prover_msg);
163            challenge = Some(transcript.get_and_append_challenge(b"Internal round")?);
164        }
165
166        end_timer!(start);
167        Ok(IOPProof {
168            proofs: prover_msgs,
169        })
170    }
171
172    /// verify the claimed sum using the proof
173    fn verify(
174        claimed_sum: F,
175        proof: &Self::Proof,
176        domain_info: &Self::DomainInfo,
177        transcript: &mut Self::Transcript,
178    ) -> Result<Self::SubClaim, PolyIOPErrors> {
179        let start = start_timer!(|| "sum check verify");
180
181        transcript.append_domain_info(domain_info)?;
182        let mut verifier_state = IOPVerifierState::verifier_init(domain_info);
183        for i in 0..domain_info.num_variables {
184            let prover_msg = proof.proofs.get(i).expect("proof is incomplete");
185            transcript.append_prover_message(prover_msg)?;
186            IOPVerifierState::verify_round_and_update_state(
187                &mut verifier_state,
188                prover_msg,
189                transcript,
190            )?;
191        }
192
193        let res = IOPVerifierState::check_and_generate_subclaim(&verifier_state, &claimed_sum);
194
195        end_timer!(start);
196        res
197    }
198}
199
200#[cfg(test)]
201mod test {
202
203    use super::*;
204    use ark_bls12_381::Fr;
205    use ark_ff::UniformRand;
206    use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
207    use ark_std::test_rng;
208    use std::rc::Rc;
209
210    fn test_sumcheck(
211        nv: usize,
212        num_multiplicands_range: (usize, usize),
213        num_products: usize,
214    ) -> Result<(), PolyIOPErrors> {
215        let mut rng = test_rng();
216        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
217
218        let (poly, asserted_sum) =
219            VirtualPolynomial::rand(nv, num_multiplicands_range, num_products, &mut rng)?;
220        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
221        let poly_info = poly.domain_info.clone();
222        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
223        let subclaim = <PolyIOP<Fr> as SumCheck<Fr>>::verify(
224            asserted_sum,
225            &proof,
226            &poly_info,
227            &mut transcript,
228        )?;
229        assert!(
230            poly.evaluate(&subclaim.point).unwrap() == subclaim.expected_evaluation,
231            "wrong subclaim"
232        );
233        Ok(())
234    }
235
236    fn test_sumcheck_internal(
237        nv: usize,
238        num_multiplicands_range: (usize, usize),
239        num_products: usize,
240    ) -> Result<(), PolyIOPErrors> {
241        let mut rng = test_rng();
242        let (poly, asserted_sum) =
243            VirtualPolynomial::<Fr>::rand(nv, num_multiplicands_range, num_products, &mut rng)?;
244        let poly_info = poly.domain_info.clone();
245        let mut prover_state = IOPProverState::prover_init(&poly)?;
246        let mut verifier_state = IOPVerifierState::verifier_init(&poly_info);
247        let mut challenge = None;
248        let mut transcript = IOPTranscript::new(b"a test transcript");
249        transcript
250            .append_message(b"testing", b"initializing transcript for testing")
251            .unwrap();
252        for _ in 0..poly.domain_info.num_variables {
253            let prover_message =
254                IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge)
255                    .unwrap();
256
257            challenge = Some(
258                IOPVerifierState::verify_round_and_update_state(
259                    &mut verifier_state,
260                    &prover_message,
261                    &mut transcript,
262                )
263                .unwrap(),
264            );
265        }
266        let subclaim =
267            IOPVerifierState::check_and_generate_subclaim(&verifier_state, &asserted_sum)
268                .expect("fail to generate subclaim");
269        assert!(
270            poly.evaluate(&subclaim.point).unwrap() == subclaim.expected_evaluation,
271            "wrong subclaim"
272        );
273        Ok(())
274    }
275
276    #[test]
277    fn test_trivial_polynomial() -> Result<(), PolyIOPErrors> {
278        let nv = 1;
279        let num_multiplicands_range = (4, 13);
280        let num_products = 5;
281
282        test_sumcheck(nv, num_multiplicands_range, num_products)?;
283        test_sumcheck_internal(nv, num_multiplicands_range, num_products)
284    }
285    #[test]
286    fn test_normal_polynomial() -> Result<(), PolyIOPErrors> {
287        let nv = 12;
288        let num_multiplicands_range = (4, 9);
289        let num_products = 5;
290
291        test_sumcheck(nv, num_multiplicands_range, num_products)?;
292        test_sumcheck_internal(nv, num_multiplicands_range, num_products)
293    }
294    #[test]
295    fn zero_polynomial_should_error() {
296        let nv = 0;
297        let num_multiplicands_range = (4, 13);
298        let num_products = 5;
299
300        assert!(test_sumcheck(nv, num_multiplicands_range, num_products).is_err());
301        assert!(test_sumcheck_internal(nv, num_multiplicands_range, num_products).is_err());
302    }
303
304    #[test]
305    fn test_extract_sum() -> Result<(), PolyIOPErrors> {
306        let mut rng = test_rng();
307        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
308        let (poly, asserted_sum) = VirtualPolynomial::<Fr>::rand(8, (3, 4), 3, &mut rng)?;
309
310        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
311        assert_eq!(
312            <PolyIOP<Fr> as SumCheck<Fr>>::extract_sum(&proof),
313            asserted_sum
314        );
315        Ok(())
316    }
317
318    #[test]
319    /// Test that the memory usage of shared-reference is linear to number of
320    /// unique MLExtensions instead of total number of multiplicands.
321    fn test_shared_reference() -> Result<(), PolyIOPErrors> {
322        let mut rng = test_rng();
323        let ml_extensions: Vec<_> = (0..5)
324            .map(|_| Rc::new(DenseMultilinearExtension::<Fr>::rand(8, &mut rng)))
325            .collect();
326        let mut poly = VirtualPolynomial::new(8);
327        poly.add_mle_list(
328            vec![
329                ml_extensions[2].clone(),
330                ml_extensions[3].clone(),
331                ml_extensions[0].clone(),
332            ],
333            Fr::rand(&mut rng),
334        )?;
335        poly.add_mle_list(
336            vec![
337                ml_extensions[1].clone(),
338                ml_extensions[4].clone(),
339                ml_extensions[4].clone(),
340            ],
341            Fr::rand(&mut rng),
342        )?;
343        poly.add_mle_list(
344            vec![
345                ml_extensions[3].clone(),
346                ml_extensions[2].clone(),
347                ml_extensions[1].clone(),
348            ],
349            Fr::rand(&mut rng),
350        )?;
351        poly.add_mle_list(
352            vec![ml_extensions[0].clone(), ml_extensions[0].clone()],
353            Fr::rand(&mut rng),
354        )?;
355        poly.add_mle_list(vec![ml_extensions[4].clone()], Fr::rand(&mut rng))?;
356
357        assert_eq!(poly.flattened_ml_extensions.len(), 5);
358
359        // test memory usage for prover
360        let prover = IOPProverState::prover_init(&poly).unwrap();
361        assert_eq!(prover.poly.flattened_ml_extensions.len(), 5);
362        drop(prover);
363
364        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
365        let poly_info = poly.domain_info.clone();
366        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
367        let asserted_sum = <PolyIOP<Fr> as SumCheck<Fr>>::extract_sum(&proof);
368
369        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
370        let subclaim = <PolyIOP<Fr> as SumCheck<Fr>>::verify(
371            asserted_sum,
372            &proof,
373            &poly_info,
374            &mut transcript,
375        )?;
376        assert!(
377            poly.evaluate(&subclaim.point)? == subclaim.expected_evaluation,
378            "wrong subclaim"
379        );
380        Ok(())
381    }
382}