MVPoly: replace signature of from_variable

From now, it will be using `Variable` from Kimchi/expr.

mvpoly/src/lib.rs

@@ -10,14 +10,12 @@
 //! "Expressions", as defined in the [kimchi] crate, can be converted into a
 //! multi-variate polynomial using the `from_expr` method.
 
-use std::collections::HashMap;
-
 use ark_ff::PrimeField;
-use kimchi::circuits::{
-    expr::{ConstantExpr, ConstantExprInner, ConstantTerm, Expr, ExprInner, Operations, Variable},
-    gate::CurrOrNext,
+use kimchi::circuits::expr::{
+    ConstantExpr, ConstantExprInner, ConstantTerm, Expr, ExprInner, Operations, Variable,
 };
 use rand::RngCore;
+use std::collections::HashMap;
 
 pub mod monomials;
 pub mod pbt;
@@ -95,7 +93,7 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
     /// used.
     /// For [crate::monomials], the output must be the index of the variable,
     /// starting from `0`.
-    fn from_variable<Column: Into<usize>>(var: Column) -> Self;
+    fn from_variable<Column: Into<usize>>(var: Variable<Column>) -> Self;
 
     fn from_constant<ChallengeTerm: Clone>(op: Operations<ConstantExprInner<F, ChallengeTerm>>) -> Self {
         use kimchi::circuits::expr::Operations::*;
@@ -160,9 +158,8 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
                         unimplemented!("Not used in this context")
                     }
                     ExprInner::Constant(c) => Self::from_constant(c),
-                    ExprInner::Cell(Variable { col, row }) => {
-                        assert_eq!(row, CurrOrNext::Curr, "Only current row is supported for now. You cannot reference the next row");
-                        Self::from_variable(col)
+                    ExprInner::Cell(var) => {
+                        Self::from_variable(var)
                     }
                 }
             }

mvpoly/src/monomials.rs

@@ -1,4 +1,5 @@
 use ark_ff::{One, PrimeField, Zero};
+use kimchi::circuits::{expr::Variable, gate::CurrOrNext};
 use num_integer::binomial;
 use rand::RngCore;
 use std::{
@@ -426,8 +427,14 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
         prime::Dense::random(rng, max_degree).into()
     }
 
-    fn from_variable<Column: Into<usize>>(var: Column) -> Self {
-        let var_usize: usize = var.into();
+    fn from_variable<Column: Into<usize>>(var: Variable<Column>) -> Self {
+        let Variable { col, row } = var;
+        assert_eq!(
+            row,
+            CurrOrNext::Curr,
+            "Only current row is supported for now. You cannot reference the next row"
+        );
+        let var_usize: usize = col.into();
         let mut monomials = HashMap::new();
         let exponents: [usize; N] = std::array::from_fn(|i| if i == var_usize { 1 } else { 0 });
         monomials.insert(exponents, F::one());

mvpoly/src/prime.rs

@@ -148,6 +148,7 @@ use std::{
 };
 
 use ark_ff::{One, PrimeField, Zero};
+use kimchi::circuits::{expr::Variable, gate::CurrOrNext};
 use num_integer::binomial;
 use o1_utils::FieldHelpers;
 use rand::{Rng, RngCore};
@@ -337,11 +338,17 @@ impl<F: PrimeField, const N: usize, const D: usize> MVPoly<F, N, D> for Dense<F,
             })
     }
 
-    fn from_variable<Column: Into<usize>>(var: Column) -> Self {
+    fn from_variable<Column: Into<usize>>(var: Variable<Column>) -> Self {
+        let Variable { col, row } = var;
+        assert_eq!(
+            row,
+            CurrOrNext::Curr,
+            "Only current row is supported for now. You cannot reference the next row"
+        );
         let mut res = Self::zero();
         let mut prime_gen = PrimeNumberGenerator::new();
         let primes = prime_gen.get_first_nth_primes(N);
-        let var_usize: usize = var.into();
+        let var_usize: usize = col.into();
         assert!(primes.contains(&var_usize), "The usize representation of the variable must be a prime number, and unique for each variable");
         let inv_var = res
             .normalized_indices

mvpoly/tests/monomials.rs

@@ -476,7 +476,10 @@ fn test_build_from_variable() {
 
     let mut rng = o1_utils::tests::make_test_rng(None);
     let idx: usize = rng.gen_range(0..4);
-    let p = Sparse::<Fp, 4, 3>::from_variable(Column::X(idx));
+    let p = Sparse::<Fp, 4, 3>::from_variable(Variable {
+        col: Column::X(idx),
+        row: CurrOrNext::Curr,
+    });
 
     let eval: [Fp; 4] = std::array::from_fn(|_i| Fp::rand(&mut rng));
 

mvpoly/tests/prime.rs

@@ -246,34 +246,6 @@ fn test_mul_by_scalar_with_from() {
     assert_eq!(result1, result2);
 }
 
-#[test]
-fn test_from_variable() {
-    // Test for y variable (index 2)
-    let y = Dense::<Fp, 4, 5>::from_variable(2_usize);
-    assert_eq!(y[1], Fp::one());
-    assert_eq!(y[0], Fp::zero());
-    assert_eq!(y[2], Fp::zero());
-    assert_eq!(y[3], Fp::zero());
-    assert_eq!(y[4], Fp::zero());
-    assert_eq!(y[5], Fp::zero());
-
-    // Test for z variable (index 3)
-    let z = Dense::<Fp, 4, 5>::from_variable(3_usize);
-    assert_eq!(z[0], Fp::zero());
-    assert_eq!(z[1], Fp::zero());
-    assert_eq!(z[2], Fp::one());
-    assert_eq!(z[3], Fp::zero());
-    assert_eq!(z[4], Fp::zero());
-
-    // Test for w variable (index 5)
-    let w = Dense::<Fp, 4, 5>::from_variable(5_usize);
-    assert_eq!(w[0], Fp::zero());
-    assert_eq!(w[1], Fp::zero());
-    assert_eq!(w[2], Fp::zero());
-    assert_eq!(w[3], Fp::zero());
-    assert_eq!(w[4], Fp::one());
-}
-
 #[test]
 fn test_from_variable_column() {
     // Simulate a real usecase
@@ -292,7 +264,10 @@ fn test_from_variable_column() {
         }
     }
 
-    let p = Dense::<Fp, 4, 5>::from_variable(Column::X(0));
+    let p = Dense::<Fp, 4, 5>::from_variable(Variable {
+        col: Column::X(0),
+        row: CurrOrNext::Curr,
+    });
     assert_eq!(p[0], Fp::zero());
     assert_eq!(p[1], Fp::one());
     assert_eq!(p[2], Fp::zero());
@@ -301,15 +276,21 @@ fn test_from_variable_column() {
     assert_eq!(p[5], Fp::zero());
 
     // Test for z variable (index 3)
-    let p = Dense::<Fp, 4, 5>::from_variable(Column::X(1));
+    let p = Dense::<Fp, 4, 5>::from_variable(Variable {
+        col: Column::X(1),
+        row: CurrOrNext::Curr,
+    });
     assert_eq!(p[0], Fp::zero());
     assert_eq!(p[1], Fp::zero());
     assert_eq!(p[2], Fp::one());
     assert_eq!(p[3], Fp::zero());
     assert_eq!(p[4], Fp::zero());
 
     // Test for w variable (index 5)
-    let p = Dense::<Fp, 4, 5>::from_variable(Column::X(2));
+    let p = Dense::<Fp, 4, 5>::from_variable(Variable {
+        col: Column::X(2),
+        row: CurrOrNext::Curr,
+    });
     assert_eq!(p[0], Fp::zero());
     assert_eq!(p[1], Fp::zero());
     assert_eq!(p[2], Fp::zero());
@@ -582,10 +563,18 @@ fn test_is_constant() {
     let p = Dense::<Fp, 4, 5>::zero();
     assert!(p.is_constant());
 
-    let p = Dense::<Fp, 4, 5>::from_variable(2_usize);
+    let p = {
+        let mut res = Dense::<Fp, 4, 5>::zero();
+        res.add_monomial([1, 0, 0, 0], Fp::one());
+        res
+    };
     assert!(!p.is_constant());
 
-    let p = Dense::<Fp, 4, 5>::from_variable(3_usize);
+    let p = {
+        let mut res = Dense::<Fp, 4, 5>::zero();
+        res.add_monomial([0, 1, 0, 0], Fp::one());
+        res
+    };
     assert!(!p.is_constant());
 
     // This might be flaky