boilerplate + fix edge cases

TomWambsgans · TomWambsgans · commit 324542940567 · 2025-10-31T00:03:44.000+04:00
diff --git a/backend/src/mle/mle_group.rs b/backend/src/mle/mle_group.rs
@@ -55,4 +55,11 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroup<'a, EF> {
             Self::Ref(_) => None,
         }
     }
+
+    pub fn is_packed(&self) -> bool {
+        match self {
+            Self::Owned(owned) => owned.is_packed(),
+            Self::Ref(r) => r.is_packed(),
+        }
+    }
 }
diff --git a/backend/src/mle/mle_single.rs b/backend/src/mle/mle_single.rs
@@ -52,4 +52,26 @@ impl<'a, EF: ExtensionField<PF<EF>>> Mle<'a, EF> {
             Self::Ref(poly) => poly.pack(),
         }
     }
+
+    pub fn unpack(&'a self) -> Self {
+        match self {
+            Self::Owned(poly) => poly.unpack(),
+            Self::Ref(poly) => poly.unpack(),
+        }
+    }
+
+    pub fn is_packed(&self) -> bool {
+        self.by_ref().is_packed()
+    }
+
+    pub fn n_vars(&self) -> usize {
+        self.by_ref().n_vars()
+    }
+
+    pub fn as_owned_or_clone(self) -> MleOwned<EF> {
+        match self {
+            Self::Owned(o) => o,
+            Self::Ref(r) => r.clone_to_owned(),
+        }
+    }
 }
diff --git a/backend/src/mle/mle_single_owned.rs b/backend/src/mle/mle_single_owned.rs
@@ -1,7 +1,7 @@
 use fiat_shamir::*;
 use p3_field::ExtensionField;
 
-use crate::{pack_extension, Mle, MleRef, MultilinearPoint};
+use crate::{Mle, MleRef, MultilinearPoint, pack_extension, unpack_extension};
 use p3_field::PackedValue;
 
 #[derive(Debug, Clone)]
@@ -112,4 +112,45 @@ impl<EF: ExtensionField<PF<EF>>> MleOwned<EF> {
             Self::ExtensionPacked(_) => Mle::Ref(self.by_ref()),
         }
     }
+
+    pub fn unpack<'a>(&'a self) -> Mle<'a, EF> {
+        match self {
+            Self::Base(v) => Mle::Ref(MleRef::Base(v)),
+            Self::Extension(v) => Mle::Ref(MleRef::Extension(v)),
+            Self::BasePacked(pb) => Mle::Ref(MleRef::Base(PFPacking::<EF>::unpack_slice(pb))),
+            Self::ExtensionPacked(ep) => Mle::Owned(MleOwned::Extension(unpack_extension(ep))),
+        }
+    }
+
+    pub fn is_packed(&self) -> bool {
+        self.by_ref().is_packed()
+    }
+
+    pub fn n_vars(&self) -> usize {
+        self.by_ref().n_vars()
+    }
+
+    pub fn halve(mut self) -> Self {
+        match &mut self {
+            Self::Base(v) => {
+                v.truncate(v.len() / 2);
+            }
+            Self::Extension(v) => {
+                v.truncate(v.len() / 2);
+            }
+            Self::BasePacked(v) => {
+                if v.len() == 1 {
+                    return self.unpack().by_ref().clone_to_owned().halve();
+                }
+                v.truncate(v.len() / 2);
+            }
+            Self::ExtensionPacked(v) => {
+                if v.len() == 1 {
+                    return self.unpack().by_ref().clone_to_owned().halve();
+                }
+                v.truncate(v.len() / 2);
+            }
+        }
+        self
+    }
 }
diff --git a/backend/src/mle/mle_single_ref.rs b/backend/src/mle/mle_single_ref.rs
@@ -86,4 +86,49 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleRef<'a, EF> {
             Self::ExtensionPacked(_) => Mle::Ref(self.clone()),
         }
     }
+
+    pub fn unpack(&self) -> Mle<'a, EF> {
+        match self {
+            Self::Base(v) => Mle::Ref(MleRef::Base(v)),
+            Self::Extension(v) => Mle::Ref(MleRef::Extension(v)),
+            Self::BasePacked(pb) => Mle::Ref(MleRef::Base(PFPacking::<EF>::unpack_slice(pb))),
+            Self::ExtensionPacked(ep) => Mle::Owned(MleOwned::Extension(unpack_extension(ep))),
+        }
+    }
+
+    pub fn pack_if(&self, cond: bool) -> Mle<'a, EF> {
+        if cond {
+            self.pack()
+        } else {
+            Mle::Ref(self.clone())
+        }
+    }
+
+    pub fn clone_to_owned(&self) -> MleOwned<EF> {
+        match self {
+            Self::Base(v) => MleOwned::Base(v.to_vec()),
+            Self::Extension(v) => MleOwned::Extension(v.to_vec()),
+            Self::BasePacked(pb) => MleOwned::BasePacked(pb.to_vec()),
+            Self::ExtensionPacked(ep) => MleOwned::ExtensionPacked(ep.to_vec()),
+        }
+    }
+
+    pub fn fold(&self, scalars: &[EF]) -> MleOwned<EF> {
+        match self {
+            Self::Base(pols) => MleOwned::Extension(fold_multilinear(pols, scalars, &|a, b| b * a)),
+            Self::Extension(pols) => {
+                MleOwned::Extension(fold_multilinear(pols, scalars, &|a, b| b * a))
+            }
+            Self::BasePacked(pols) => {
+                let scalars_packed = scalars
+                    .iter()
+                    .map(|&s| EFPacking::<EF>::from(s))
+                    .collect::<Vec<_>>();
+                MleOwned::ExtensionPacked(fold_multilinear(pols, &scalars_packed, &|a, b| b * a))
+            }
+            Self::ExtensionPacked(pols) => {
+                MleOwned::ExtensionPacked(fold_multilinear(pols, scalars, &|a, b| a * b))
+            }
+        }
+    }
 }
diff --git a/sumcheck/src/product_computation.rs b/sumcheck/src/product_computation.rs
@@ -61,6 +61,9 @@ pub fn run_product_sumcheck<EF: ExtensionField<PF<EF>>>(
                 EFPacking::<EF>::to_ext_iter([e]).collect()
             })
         }
+        (MleRef::Base(evals), MleRef::Extension(weights)) => {
+            compute_product_sumcheck_polynomial(evals, &weights, sum, |e| vec![e])
+        }
         _ => unimplemented!(),
     };
 
@@ -69,20 +72,34 @@ pub fn run_product_sumcheck<EF: ExtensionField<PF<EF>>>(
     let r1: EF = prover_state.sample();
     sum = first_sumcheck_poly.evaluate(r1);
 
-    if n_rounds < 2 {
-        unimplemented!()
+    if n_rounds == 1 {
+        return (
+            MultilinearPoint(vec![r1]),
+            sum,
+            pol_a.fold(&[EF::ONE - r1, r1]),
+            pol_b.fold(&[EF::ONE - r1, r1]),
+        );
     }
 
     let (second_sumcheck_poly, folded) = match (pol_a, pol_b) {
         (MleRef::BasePacked(evals), MleRef::ExtensionPacked(weights)) => {
-            fold_and_compute_product_sumcheck_polynomial(&evals, &weights, r1, sum, |e| {
-                EFPacking::<EF>::to_ext_iter([e]).collect()
-            })
+            let (second_sumcheck_poly, folded) =
+                fold_and_compute_product_sumcheck_polynomial(&evals, &weights, r1, sum, |e| {
+                    EFPacking::<EF>::to_ext_iter([e]).collect()
+                });
+            (second_sumcheck_poly, MleGroupOwned::ExtensionPacked(folded))
         }
         (MleRef::ExtensionPacked(evals), MleRef::ExtensionPacked(weights)) => {
-            fold_and_compute_product_sumcheck_polynomial(evals, &weights, r1, sum, |e| {
-                EFPacking::<EF>::to_ext_iter([e]).collect()
-            })
+            let (second_sumcheck_poly, folded) =
+                fold_and_compute_product_sumcheck_polynomial(evals, &weights, r1, sum, |e| {
+                    EFPacking::<EF>::to_ext_iter([e]).collect()
+                });
+            (second_sumcheck_poly, MleGroupOwned::ExtensionPacked(folded))
+        }
+        (MleRef::Base(evals), MleRef::Extension(weights)) => {
+            let (second_sumcheck_poly, folded) =
+                fold_and_compute_product_sumcheck_polynomial(evals, &weights, r1, sum, |e| vec![e]);
+            (second_sumcheck_poly, MleGroupOwned::Extension(folded))
         }
         _ => unimplemented!(),
     };
@@ -94,7 +111,7 @@ pub fn run_product_sumcheck<EF: ExtensionField<PF<EF>>>(
 
     let (mut challenges, folds, sum) = sumcheck_prove_many_rounds(
         1,
-        MleGroupOwned::ExtensionPacked(folded),
+        folded,
         Some(vec![EF::ONE - r2, r2]),
         &ProductComputation,
         &[],
diff --git a/sumcheck/src/prove.rs b/sumcheck/src/prove.rs
@@ -127,17 +127,18 @@ where
             });
             (eq_point, eq_mle)
         });
+
     let mut n_vars = multilinears.by_ref().n_vars();
     if let Some(prev_folding_factors) = &prev_folding_factors {
         n_vars -= log2_strict_usize(prev_folding_factors.len());
     }
     if let Some((eq_point, eq_mle)) = &eq_factor {
         assert_eq!(eq_point.len(), n_vars - skip + 1);
         assert_eq!(eq_mle.by_ref().n_vars(), eq_point.len() - 1);
-        assert_eq!(
-            eq_mle.by_ref().is_packed(),
-            multilinears.by_ref().is_packed()
-        );
+        if eq_mle.by_ref().is_packed() && !multilinears.is_packed() {
+            assert!(eq_point.len() < packing_log_width::<EF>());
+            multilinears = multilinears.by_ref().unpack().into();
+        }
     }
 
     let mut challenges = Vec::new();
@@ -147,9 +148,7 @@ where
             // unpack
             multilinears = multilinears.by_ref().unpack().into();
             if let Some((_, eq_mle)) = &mut eq_factor {
-                *eq_mle = MleOwned::Extension(unpack_extension(
-                    eq_mle.by_ref().as_extension_packed().unwrap(),
-                ));
+                *eq_mle = eq_mle.by_ref().unpack().as_owned_or_clone();
             }
         }
 

Original file line number	Diff line number	Diff line change
`@@ -55,4 +55,11 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroup<'a, EF> {`
`55`	`55`	`Self::Ref(_) => None,`
`56`	`56`	`}`
`57`	`57`	`}`
	`58`	`+`
	`59`	`+ pub fn is_packed(&self) -> bool {`
	`60`	`+ match self {`
	`61`	`+ Self::Owned(owned) => owned.is_packed(),`
	`62`	`+ Self::Ref(r) => r.is_packed(),`
	`63`	`+ }`
	`64`	`+ }`
`58`	`65`	`}`