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