Proof size reduction for sumchecks (see https://eprint.iacr.org/2024/108.pdf, section 3)

Author: TomWambsgans

README.md

@@ -36,8 +36,8 @@ cargo run --release -- xmss --n-signatures 1400
 
 | WHIR rate \ regime | Proven               | Conjectured          |
 | ------------------ | -------------------- | -------------------- |
-| 1/2                | 775 XMSS/s - 374 KiB | 775 XMSS/s - 204 KiB |
-| 1/4                | 650 XMSS/s - 246 KiB | 650 XMSS/s - 146 KiB |
+| 1/2                | 800 XMSS/s - 355 KiB | 800 XMSS/s - 188 KiB |
+| 1/4                | 700 XMSS/s - 229 KiB | 650 XMSS/s - 130 KiB |
 
 (Proving throughput - proof size)
 
@@ -52,7 +52,7 @@ cargo run --release -- recursion --n 2
 
 | Proven          | Conjectured     |
 | --------------- | --------------- |
-| 0.77s - 223 KiB | 0.59s - 128 KiB |
+| 0.75s - 188 KiB | 0.57s - 116 KiB |
 
 
 ### Bonus: unbounded recursive aggregation

crates/air/src/verify.rs

@@ -22,15 +22,14 @@ where
         .unwrap_or_else(|| verifier_state.sample_vec(log_n_rows));
     assert_eq!(zerocheck_challenges.len(), log_n_rows);
 
-    let (sc_sum, outer_statement) = sumcheck_verify::<EF>(verifier_state, log_n_rows, air.degree_air() + 1)?;
-    if sc_sum
-        != virtual_column_statement
-            .as_ref()
-            .map(|st| st.value)
-            .unwrap_or_else(|| EF::ZERO)
-    {
-        return Err(ProofError::InvalidProof);
-    }
+    let expected_sum = virtual_column_statement.as_ref().map(|st| st.value).unwrap_or(EF::ZERO);
+    let outer_statement = sumcheck_verify(
+        verifier_state,
+        log_n_rows,
+        air.degree_air() + 1, // +1 for the eq factor
+        expected_sum,
+        Some(&zerocheck_challenges),
+    )?;
 
     let inner_evals = verifier_state.next_extension_scalars_vec(air.n_columns() + air.n_down_columns())?;
 
@@ -74,11 +73,7 @@ fn open_columns<A: Air, EF: ExtensionField<PF<EF>>>(
 
     let inner_sum: EF = dot_product(evals_down.into_iter(), batching_scalar_powers.iter().copied());
 
-    let (inner_sum_retrieved, inner_sumcheck_stement) = sumcheck_verify(verifier_state, log_n_rows, 2)?;
-
-    if inner_sum != inner_sum_retrieved {
-        return Err(ProofError::InvalidProof);
-    }
+    let inner_sumcheck_stement = sumcheck_verify(verifier_state, log_n_rows, 2, inner_sum, None)?;
 
     let matrix_down_sc_eval = next_mle(outer_sumcheck_challenge, &inner_sumcheck_stement.point);
 

crates/backend/fiat-shamir/src/challenger.rs

@@ -36,6 +36,16 @@ impl<F: PrimeField64, P: Compression<[F; WIDTH]>> Challenger<F, P> {
         self.state = self.hash_state_with(&value);
     }
 
+    pub fn observe_scalars(&mut self, scalars: &[F]) {
+        for chunk in scalars.chunks(RATE) {
+            let mut buffer = [F::ZERO; RATE];
+            for (i, val) in chunk.iter().enumerate() {
+                buffer[i] = *val;
+            }
+            self.observe(buffer);
+        }
+    }
+
     pub fn sample_many(&mut self, n: usize) -> Vec<[F; RATE]> {
         let mut sampled = Vec::with_capacity(n);
         for i in 0..n + 1 {

crates/backend/fiat-shamir/src/lib.rs

@@ -4,9 +4,6 @@ pub use errors::*;
 mod prover;
 pub use prover::*;
 
-mod verifier;
-pub use verifier::*;
-
 mod utils;
 pub use utils::*;
 
@@ -16,7 +13,10 @@ mod traits;
 pub use traits::*;
 
 mod transcript;
-pub use transcript::*;
+pub use transcript::{DIGEST_LEN_FE, MerkleOpening, MerklePath, MerklePaths, Proof, RawProof};
 
 mod merkle_pruning;
 pub(crate) use merkle_pruning::*;
+
+mod verifier;
+pub use verifier::*;

crates/backend/fiat-shamir/src/prover.rs

@@ -1,4 +1,5 @@
 use crate::{
+    MerklePaths, PrunedMerklePaths,
     challenger::{Challenger, RATE, WIDTH},
     *,
 };
@@ -14,7 +15,8 @@ use symetric::Compression;
 #[derive(Debug)]
 pub struct ProverState<EF: ExtensionField<PF<EF>>, P> {
     challenger: Challenger<PF<EF>, P>,
-    transcript: Proof<PF<EF>>,
+    transcript: Vec<PF<EF>>,
+    merkle_paths: Vec<PrunedMerklePaths<PF<EF>, PF<EF>>>,
 }
 
 impl<EF: ExtensionField<PF<EF>>, P: Compression<[PF<EF>; WIDTH]>> ProverState<EF, P>
@@ -26,20 +28,16 @@ where
         assert!(EF::DIMENSION <= RATE);
         Self {
             challenger: Challenger::new(compressor),
-            transcript: Proof(Vec::new()),
+            transcript: Vec::new(),
+            merkle_paths: Vec::new(),
         }
     }
 
-    pub fn raw_proof(self) -> RawProof<PF<EF>> {
-        self.transcript.into_raw_proof()
-    }
-
-    pub fn pruned_proof(&self) -> PrunedProof<PF<EF>> {
-        self.transcript.clone().prune()
-    }
-
-    pub fn into_pruned_proof(self) -> PrunedProof<PF<EF>> {
-        self.transcript.prune()
+    pub fn into_proof(self) -> Proof<PF<EF>> {
+        Proof {
+            transcript: self.transcript,
+            merkle_paths: self.merkle_paths,
+        }
     }
 }
 
@@ -62,14 +60,8 @@ where
     PF<EF>: PrimeField64,
 {
     fn add_base_scalars(&mut self, scalars: &[PF<EF>]) {
-        self.transcript.0.push(TranscriptData::Interraction(scalars.to_vec()));
-        for chunk in scalars.chunks(RATE) {
-            let mut buffer = [PF::<EF>::ZERO; RATE];
-            for (i, val) in chunk.iter().enumerate() {
-                buffer[i] = *val;
-            }
-            self.challenger.observe(buffer);
-        }
+        self.challenger.observe_scalars(scalars);
+        self.transcript.extend_from_slice(scalars);
     }
 
     fn state(&self) -> String {
@@ -81,12 +73,28 @@ where
                 .map(|f| f.to_string())
                 .collect::<Vec<_>>()
                 .join(", "),
-            self.transcript.0.len()
+            self.transcript.len()
         )
     }
 
+    fn add_sumcheck_polynomial(&mut self, coeffs: &[EF], eq_alpha: Option<EF>) {
+        match eq_alpha {
+            None => {
+                let scalars = flatten_scalars_to_base(coeffs);
+                self.challenger.observe_scalars(&scalars);
+                self.transcript.extend_from_slice(&scalars[EF::DIMENSION..]); // c0 reconstructed by verifier from claimed_sum
+            }
+            Some(alpha) => {
+                let bare_scalars = flatten_scalars_to_base(coeffs);
+                let full_scalars = flatten_scalars_to_base(&expand_bare_to_full(coeffs, alpha));
+                self.challenger.observe_scalars(&full_scalars);
+                self.transcript.extend_from_slice(&bare_scalars[EF::DIMENSION..]); // h0 reconstructed by verifier from claimed_sum
+            }
+        }
+    }
+
     fn hint_merkle_paths_base(&mut self, paths: Vec<MerklePath<PF<EF>, PF<EF>>>) {
-        self.transcript.0.push(TranscriptData::MerklePaths(MerklePaths(paths)));
+        self.merkle_paths.push(MerklePaths(paths).prune());
     }
 
     fn pow_grinding(&mut self, bits: usize) {
@@ -134,14 +142,10 @@ where
             })
             .expect("failed to find witness");
 
-        let witness_found = witness_found.lock().unwrap().unwrap();
+        let witness = witness_found.lock().unwrap().unwrap();
 
-        self.challenger.observe({
-            let mut value = [PF::<EF>::ZERO; RATE];
-            value[0] = witness_found;
-            value
-        });
+        self.challenger.observe_scalars(&[witness]);
         assert!(self.challenger.state[0].as_canonical_u64() & ((1 << bits) - 1) == 0);
-        self.transcript.0.push(TranscriptData::GrindingWitness(witness_found));
+        self.transcript.push(witness);
     }
 }

crates/backend/fiat-shamir/src/traits.rs

@@ -1,6 +1,8 @@
 use field::ExtensionField;
 
-use crate::{MerkleOpening, MerklePath, PF, ProofError, flatten_scalars_to_base, pack_scalars_to_extension};
+use crate::{
+    MerkleOpening, MerklePath, PF, ProofError, ProofResult, flatten_scalars_to_base, pack_scalars_to_extension,
+};
 
 pub trait ChallengeSampler<EF> {
     fn sample_vec(&mut self, len: usize) -> Vec<EF>;
@@ -15,6 +17,7 @@ pub trait FSProver<EF: ExtensionField<PF<EF>>>: ChallengeSampler<EF> {
     fn add_base_scalars(&mut self, scalars: &[PF<EF>]);
     fn pow_grinding(&mut self, bits: usize);
     fn hint_merkle_paths_base(&mut self, paths: Vec<MerklePath<PF<EF>, PF<EF>>>);
+    fn add_sumcheck_polynomial(&mut self, coeffs: &[EF], eq_alpha: Option<EF>);
 
     fn add_extension_scalars(&mut self, scalars: &[EF]) {
         self.add_base_scalars(&flatten_scalars_to_base(scalars));
@@ -43,6 +46,12 @@ pub trait FSVerifier<EF: ExtensionField<PF<EF>>>: ChallengeSampler<EF> {
     fn next_base_scalars_vec(&mut self, n: usize) -> Result<Vec<PF<EF>>, ProofError>;
     fn next_merkle_opening(&mut self) -> Result<MerkleOpening<PF<EF>>, ProofError>;
     fn check_pow_grinding(&mut self, bits: usize) -> Result<(), ProofError>;
+    fn next_sumcheck_polynomial(
+        &mut self,
+        n_coeffs: usize,
+        claimed_sum: EF,
+        eq_alpha: Option<EF>,
+    ) -> ProofResult<Vec<EF>>;
 
     fn next_extension_scalars_vec(&mut self, n: usize) -> Result<Vec<EF>, ProofError> {
         Ok(pack_scalars_to_extension(