sumcheck refactor & unittest to demonstrate various num_var within single sumcheck (#862)

This PR addressed few question raised from @Jiangkm3 when deal with mpcs
poly batching tasks

### change scope
- [x] refactor `virtual_polys` out from `ceno_zkvm` to mle crate so
could be share for other crates
- [x] unittest to show sumcheck with different variables

Author: hero78119

ceno_zkvm/src/lib.rs

@@ -21,7 +21,6 @@ pub mod stats;
 pub mod structs;
 mod uint;
 mod utils;
-mod virtual_polys;
 mod witness;
 
 pub use structs::ROMType;

ceno_zkvm/src/scheme/prover.rs

@@ -10,6 +10,7 @@ use multilinear_extensions::{
     mle::{IntoMLE, MultilinearExtension},
     util::ceil_log2,
     virtual_poly::{ArcMultilinearExtension, build_eq_x_r_vec},
+    virtual_polys::VirtualPolynomials,
 };
 use p3_field::PrimeCharacteristicRing;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
@@ -33,8 +34,7 @@ use crate::{
     structs::{
         Point, ProvingKey, TowerProofs, TowerProver, TowerProverSpec, ZKVMProvingKey, ZKVMWitnesses,
     },
-    utils::{get_challenge_pows, optimal_sumcheck_threads},
-    virtual_polys::VirtualPolynomials,
+    utils::{add_mle_list_by_expr, get_challenge_pows, optimal_sumcheck_threads},
 };
 
 use super::{PublicValues, ZKVMOpcodeProof, ZKVMProof, ZKVMTableProof};
@@ -572,7 +572,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                     }
                 }
 
-                distrinct_zerocheck_terms_set.extend(virtual_polys.add_mle_list_by_expr(
+                distrinct_zerocheck_terms_set.extend(add_mle_list_by_expr(
+                    &mut virtual_polys,
                     sel_non_lc_zero_sumcheck.as_ref(),
                     witnesses.iter().collect_vec(),
                     expr,

ceno_zkvm/src/utils.rs

@@ -1,16 +1,22 @@
 use std::{
-    collections::HashMap,
+    collections::{BTreeSet, HashMap},
     fmt::Display,
     hash::Hash,
+    mem,
     panic::{self, PanicHookInfo},
 };
 
 use ff_ext::{ExtensionField, SmallField};
 use itertools::Itertools;
-use multilinear_extensions::util::max_usable_threads;
+use multilinear_extensions::{
+    util::max_usable_threads, virtual_poly::ArcMultilinearExtension,
+    virtual_polys::VirtualPolynomials,
+};
 use p3_field::Field;
 use transcript::Transcript;
 
+use crate::expression::Expression;
+
 pub fn i64_to_base<F: SmallField>(x: i64) -> F {
     if x >= 0 {
         F::from_u64(x as u64)
@@ -167,21 +173,6 @@ pub fn eval_wellform_address_vec<E: ExtensionField>(offset: u64, scaled: u64, r:
                 .sum::<E>()
 }
 
-/// transpose 2d vector without clone
-pub fn transpose<T>(v: Vec<Vec<T>>) -> Vec<Vec<T>> {
-    assert!(!v.is_empty());
-    let len = v[0].len();
-    let mut iters: Vec<_> = v.into_iter().map(|n| n.into_iter()).collect();
-    (0..len)
-        .map(|_| {
-            iters
-                .iter_mut()
-                .map(|n| n.next().unwrap())
-                .collect::<Vec<T>>()
-        })
-        .collect()
-}
-
 pub fn display_hashmap<K: Display, V: Display>(map: &HashMap<K, V>) -> String {
     format!(
         "[{}]",
@@ -223,3 +214,137 @@ where
 
     result
 }
+
+/// add mle terms into virtual poly by expression
+/// return distinct witin in set
+pub fn add_mle_list_by_expr<'a, E: ExtensionField>(
+    virtual_polys: &mut VirtualPolynomials<'a, E>,
+    selector: Option<&'a ArcMultilinearExtension<'a, E>>,
+    wit_ins: Vec<&'a ArcMultilinearExtension<'a, E>>,
+    expr: &Expression<E>,
+    challenges: &[E],
+    // sumcheck batch challenge
+    alpha: E,
+) -> BTreeSet<u16> {
+    assert!(expr.is_monomial_form());
+    let monomial_terms = expr.evaluate(
+        &|_| unreachable!(),
+        &|witness_id| vec![(E::ONE, { vec![witness_id] })],
+        &|structural_witness_id, _, _, _| vec![(E::ONE, { vec![structural_witness_id] })],
+        &|scalar| vec![(E::from(scalar), { vec![] })],
+        &|challenge_id, pow, scalar, offset| {
+            let challenge = challenges[challenge_id as usize];
+            vec![(challenge.exp_u64(pow as u64) * scalar + offset, vec![])]
+        },
+        &|mut a, b| {
+            a.extend(b);
+            a
+        },
+        &|mut a, mut b| {
+            assert!(a.len() <= 2);
+            assert!(b.len() <= 2);
+            // special logic to deal with scaledsum
+            // scaledsum second parameter must be 0
+            if a.len() == 2 {
+                assert!((a[1].0, a[1].1.is_empty()) == (E::ZERO, true));
+                a.truncate(1);
+            }
+            if b.len() == 2 {
+                assert!((b[1].0, b[1].1.is_empty()) == (E::ZERO, true));
+                b.truncate(1);
+            }
+
+            a[0].1.extend(mem::take(&mut b[0].1));
+            // return [ab]
+            vec![(a[0].0 * b[0].0, mem::take(&mut a[0].1))]
+        },
+        &|mut x, a, b| {
+            assert!(a.len() == 1 && a[0].1.is_empty()); // for challenge or constant, term should be empty
+            assert!(b.len() == 1 && b[0].1.is_empty()); // for challenge or constant, term should be empty
+            assert!(x.len() == 1 && (x[0].0, x[0].1.len()) == (E::ONE, 1)); // witin size only 1
+            if b[0].0 == E::ZERO {
+                // only include first term if b = 0
+                vec![(a[0].0, mem::take(&mut x[0].1))]
+            } else {
+                // return [ax, b]
+                vec![(a[0].0, mem::take(&mut x[0].1)), (b[0].0, vec![])]
+            }
+        },
+    );
+    for (constant, monomial_term) in monomial_terms.iter() {
+        if *constant != E::ZERO && monomial_term.is_empty() && selector.is_none() {
+            todo!("make virtual poly support pure constant")
+        }
+        let sel = selector.map(|sel| vec![sel]).unwrap_or_default();
+        let terms_polys = monomial_term
+            .iter()
+            .map(|wit_id| wit_ins[*wit_id as usize])
+            .collect_vec();
+
+        virtual_polys.add_mle_list([sel, terms_polys].concat(), *constant * alpha);
+    }
+
+    monomial_terms
+        .into_iter()
+        .flat_map(|(_, monomial_term)| monomial_term.into_iter().collect_vec())
+        .collect::<BTreeSet<u16>>()
+}
+
+#[cfg(test)]
+mod tests {
+    use ff_ext::GoldilocksExt2;
+    use itertools::Itertools;
+    use multilinear_extensions::{
+        mle::IntoMLE, virtual_poly::ArcMultilinearExtension, virtual_polys::VirtualPolynomials,
+    };
+    use p3_field::PrimeCharacteristicRing;
+
+    use crate::{
+        circuit_builder::{CircuitBuilder, ConstraintSystem},
+        expression::{Expression, ToExpr},
+        utils::add_mle_list_by_expr,
+    };
+
+    #[test]
+    fn test_add_mle_list_by_expr() {
+        type E = ff_ext::GoldilocksExt2;
+        type F = p3_goldilocks::Goldilocks;
+        let mut cs = ConstraintSystem::new(|| "test_root");
+        let mut cb = CircuitBuilder::<E>::new(&mut cs);
+        let x = cb.create_witin(|| "x");
+        let y = cb.create_witin(|| "y");
+
+        let wits_in: Vec<ArcMultilinearExtension<E>> = (0..cs.num_witin as usize)
+            .map(|_| vec![F::from_u64(1)].into_mle().into())
+            .collect();
+
+        let mut virtual_polys = VirtualPolynomials::new(1, 0);
+
+        // 3xy + 2y
+        let expr: Expression<E> = 3 * x.expr() * y.expr() + 2 * y.expr();
+
+        let distrinct_zerocheck_terms_set = add_mle_list_by_expr(
+            &mut virtual_polys,
+            None,
+            wits_in.iter().collect_vec(),
+            &expr,
+            &[],
+            GoldilocksExt2::ONE,
+        );
+        assert!(distrinct_zerocheck_terms_set.len() == 2);
+        assert!(virtual_polys.degree() == 2);
+
+        // 3x^3
+        let expr: Expression<E> = 3 * x.expr() * x.expr() * x.expr();
+        let distrinct_zerocheck_terms_set = add_mle_list_by_expr(
+            &mut virtual_polys,
+            None,
+            wits_in.iter().collect_vec(),
+            &expr,
+            &[],
+            GoldilocksExt2::ONE,
+        );
+        assert!(distrinct_zerocheck_terms_set.len() == 1);
+        assert!(virtual_polys.degree() == 3);
+    }
+}