refactor: make limbs generic to accept limb size

core/src/air/word.rs

@@ -1,5 +1,5 @@
 use std::array::IntoIter;
-use std::mem::size_of;
+
 use std::ops::{Index, IndexMut};
 
 use core::borrow::{Borrow, BorrowMut};

core/src/memory/columns.rs

@@ -1,6 +1,5 @@
 use core::borrow::{Borrow, BorrowMut};
 use sp1_derive::AlignedBorrow;
-use std::mem::size_of;
 
 use crate::air::Word;
 

core/src/operations/add.rs

@@ -3,7 +3,6 @@ use core::borrow::BorrowMut;
 use p3_air::AirBuilder;
 use p3_field::Field;
 use sp1_derive::AlignedBorrow;
-use std::mem::size_of;
 
 use crate::air::SP1AirBuilder;
 use crate::air::Word;

core/src/operations/add4.rs

@@ -5,7 +5,6 @@ use p3_air::AirBuilder;
 use p3_field::AbstractField;
 use p3_field::Field;
 use sp1_derive::AlignedBorrow;
-use std::mem::size_of;
 
 use crate::air::SP1AirBuilder;
 use crate::air::Word;

core/src/operations/add5.rs

@@ -3,7 +3,6 @@ use core::borrow::BorrowMut;
 use p3_air::AirBuilder;
 use p3_field::Field;
 use sp1_derive::AlignedBorrow;
-use std::mem::size_of;
 
 use crate::air::SP1AirBuilder;
 use crate::air::Word;

core/src/operations/and.rs

@@ -3,7 +3,6 @@ use core::borrow::BorrowMut;
 use p3_field::AbstractField;
 use p3_field::Field;
 use sp1_derive::AlignedBorrow;
-use std::mem::size_of;
 
 use crate::air::SP1AirBuilder;
 use crate::air::Word;

core/src/operations/field/field_den.rs

@@ -1,12 +1,10 @@
 use super::params::Limbs;
-use super::params::NUM_WITNESS_LIMBS;
 use super::util::{compute_root_quotient_and_shift, split_u16_limbs_to_u8_limbs};
 use super::util_air::eval_field_operation;
 use crate::air::Polynomial;
 use crate::air::SP1AirBuilder;
-use crate::utils::ec::field::FieldParameters;
-use core::borrow::{Borrow, BorrowMut};
-use core::mem::size_of;
+use crate::utils::ec::field::{limbs_from_vec, FieldParameters};
+use core::borrow::Borrow;
 use num::BigUint;
 use p3_field::PrimeField32;
 use sp1_derive::AlignedBorrow;
@@ -21,15 +19,15 @@ use std::fmt::Debug;
 /// or made generic in the future.
 #[derive(Debug, Clone, AlignedBorrow)]
 #[repr(C)]
-pub struct FieldDenCols<T> {
+pub struct FieldDenCols<T, const N: usize, const M: usize> {
     /// The result of `a den b`, where a, b are field elements
-    pub result: Limbs<T>,
-    pub(crate) carry: Limbs<T>,
-    pub(crate) witness_low: [T; NUM_WITNESS_LIMBS],
-    pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
+    pub result: Limbs<T, N>,
+    pub(crate) carry: Limbs<T, N>,
+    pub(crate) witness_low: [T; M],
+    pub(crate) witness_high: [T; M],
 }
 
-impl<F: PrimeField32> FieldDenCols<F> {
+impl<F: PrimeField32, const N: usize, const M: usize> FieldDenCols<F, N, M> {
     pub fn populate<P: FieldParameters>(
         &mut self,
         a: &BigUint,
@@ -55,11 +53,12 @@ impl<F: PrimeField32> FieldDenCols<F> {
         debug_assert!(carry < p);
         debug_assert_eq!(&carry * &p, &equation_lhs - &equation_rhs);
 
-        let p_a: Polynomial<F> = P::to_limbs_field::<F>(a).into();
-        let p_b: Polynomial<F> = P::to_limbs_field::<F>(b).into();
-        let p_p: Polynomial<F> = P::to_limbs_field::<F>(&p).into();
-        let p_result: Polynomial<F> = P::to_limbs_field::<F>(&result).into();
-        let p_carry: Polynomial<F> = P::to_limbs_field::<F>(&carry).into();
+        let p_a: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(a)).into();
+        let p_b: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(b)).into();
+        let p_p: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(&p)).into();
+        let p_result: Polynomial<F> =
+            limbs_from_vec::<F, N>(P::to_limbs_field::<F>(&result)).into();
+        let p_carry: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(&carry)).into();
 
         // Compute the vanishing polynomial.
         let vanishing_poly = if sign {
@@ -85,13 +84,13 @@ impl<F: PrimeField32> FieldDenCols<F> {
     }
 }
 
-impl<V: Copy> FieldDenCols<V> {
+impl<V: Copy, const N: usize, const M: usize> FieldDenCols<V, N, M> {
     #[allow(unused_variables)]
     pub fn eval<AB: SP1AirBuilder<Var = V>, P: FieldParameters>(
         &self,
         builder: &mut AB,
-        a: &Limbs<AB::Var>,
-        b: &Limbs<AB::Var>,
+        a: &Limbs<AB::Var, N>,
+        b: &Limbs<AB::Var, N>,
         sign: bool,
     ) where
         V: Into<AB::Expr>,
@@ -120,7 +119,7 @@ impl<V: Copy> FieldDenCols<V> {
         let p_witness_low = self.witness_low.iter().into();
         let p_witness_high = self.witness_high.iter().into();
 
-        eval_field_operation::<AB, P>(builder, &p_vanishing, &p_witness_low, &p_witness_high);
+        eval_field_operation::<AB, N, P>(builder, &p_vanishing, &p_witness_low, &p_witness_high);
     }
 }
 
@@ -135,7 +134,7 @@ mod tests {
     use crate::air::MachineAir;
 
     use crate::utils::ec::edwards::ed25519::Ed25519BaseField;
-    use crate::utils::ec::field::FieldParameters;
+    use crate::utils::ec::field::{limbs_from_vec, FieldParameters};
     use crate::utils::{uni_stark_prove as prove, uni_stark_verify as verify};
     use crate::utils::{BabyBearPoseidon2, StarkUtils};
     use crate::{air::SP1AirBuilder, runtime::ExecutionRecord};
@@ -148,14 +147,18 @@ mod tests {
     use p3_matrix::MatrixRowSlices;
     use rand::thread_rng;
     use sp1_derive::AlignedBorrow;
-    #[derive(AlignedBorrow, Debug, Clone)]
-    pub struct TestCols<T> {
-        pub a: Limbs<T>,
-        pub b: Limbs<T>,
-        pub a_den_b: FieldDenCols<T>,
+
+    #[derive(Debug, Clone, AlignedBorrow)]
+    pub struct TestCols<T, const N: usize, const M: usize> {
+        pub a: Limbs<T, N>,
+        pub b: Limbs<T, N>,
+        pub a_den_b: FieldDenCols<T, N, M>,
     }
 
-    pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8>>();
+    const NUM_LIMBS: usize = 32;
+    const NUM_WITNESS_LIMBS: usize = NUM_LIMBS * 2 - 2;
+
+    pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8, NUM_LIMBS, NUM_WITNESS_LIMBS>>();
 
     struct FieldDenChip<P: FieldParameters> {
         pub sign: bool,
@@ -205,9 +208,10 @@ mod tests {
                 .iter()
                 .map(|(a, b)| {
                     let mut row = [F::zero(); NUM_TEST_COLS];
-                    let cols: &mut TestCols<F> = row.as_mut_slice().borrow_mut();
-                    cols.a = P::to_limbs_field::<F>(a);
-                    cols.b = P::to_limbs_field::<F>(b);
+                    let cols: &mut TestCols<F, NUM_LIMBS, NUM_WITNESS_LIMBS> =
+                        row.as_mut_slice().borrow_mut();
+                    cols.a = limbs_from_vec::<F, NUM_LIMBS>(P::to_limbs_field::<F>(a));
+                    cols.b = limbs_from_vec::<F, NUM_LIMBS>(P::to_limbs_field::<F>(b));
                     cols.a_den_b.populate::<P>(a, b, self.sign);
                     row
                 })
@@ -235,7 +239,8 @@ mod tests {
     {
         fn eval(&self, builder: &mut AB) {
             let main = builder.main();
-            let local: &TestCols<AB::Var> = main.row_slice(0).borrow();
+            let local: &TestCols<AB::Var, NUM_LIMBS, NUM_WITNESS_LIMBS> =
+                main.row_slice(0).borrow();
             local
                 .a_den_b
                 .eval::<AB, P>(builder, &local.a, &local.b, self.sign);
@@ -257,7 +262,7 @@ mod tests {
     }
 
     #[test]
-    fn prove_babybear() {
+    fn prove_field_den_babybear() {
         let config = BabyBearPoseidon2::new();
         let mut challenger = config.challenger();
 

core/src/operations/field/field_inner_product.rs

@@ -1,12 +1,10 @@
 use super::params::Limbs;
-use super::params::NUM_WITNESS_LIMBS;
 use super::util::{compute_root_quotient_and_shift, split_u16_limbs_to_u8_limbs};
 use super::util_air::eval_field_operation;
 use crate::air::Polynomial;
 use crate::air::SP1AirBuilder;
-use crate::utils::ec::field::FieldParameters;
-use core::borrow::{Borrow, BorrowMut};
-use core::mem::size_of;
+use crate::utils::ec::field::{limbs_from_vec, FieldParameters};
+use core::borrow::Borrow;
 use num::BigUint;
 use num::Zero;
 use p3_field::{AbstractField, PrimeField32};
@@ -18,20 +16,24 @@ use std::fmt::Debug;
 /// or made generic in the future.
 #[derive(Debug, Clone, AlignedBorrow)]
 #[repr(C)]
-pub struct FieldInnerProductCols<T> {
+pub struct FieldInnerProductCols<T, const N: usize, const M: usize> {
     /// The result of `a inner product b`, where a, b are field elements
-    pub result: Limbs<T>,
-    pub(crate) carry: Limbs<T>,
-    pub(crate) witness_low: [T; NUM_WITNESS_LIMBS],
-    pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
+    pub result: Limbs<T, N>,
+    pub(crate) carry: Limbs<T, N>,
+    pub(crate) witness_low: [T; M],
+    pub(crate) witness_high: [T; M],
 }
 
-impl<F: PrimeField32> FieldInnerProductCols<F> {
+impl<F: PrimeField32, const N: usize, const M: usize> FieldInnerProductCols<F, N, M> {
     pub fn populate<P: FieldParameters>(&mut self, a: &[BigUint], b: &[BigUint]) -> BigUint {
-        let p_a_vec: Vec<Polynomial<F>> =
-            a.iter().map(|x| P::to_limbs_field::<F>(x).into()).collect();
-        let p_b_vec: Vec<Polynomial<F>> =
-            b.iter().map(|x| P::to_limbs_field::<F>(x).into()).collect();
+        let p_a_vec: Vec<Polynomial<F>> = a
+            .iter()
+            .map(|x| limbs_from_vec::<F, N>(P::to_limbs_field::<F>(x)).into())
+            .collect();
+        let p_b_vec: Vec<Polynomial<F>> = b
+            .iter()
+            .map(|x| limbs_from_vec::<F, N>(P::to_limbs_field::<F>(x)).into())
+            .collect();
 
         let modulus = &P::modulus();
         let inner_product = a
@@ -45,9 +47,10 @@ impl<F: PrimeField32> FieldInnerProductCols<F> {
         assert!(carry < &(2u32 * modulus));
         assert_eq!(carry * modulus, inner_product - result);
 
-        let p_modulus: Polynomial<F> = P::to_limbs_field::<F>(modulus).into();
-        let p_result: Polynomial<F> = P::to_limbs_field::<F>(result).into();
-        let p_carry: Polynomial<F> = P::to_limbs_field::<F>(carry).into();
+        let p_modulus: Polynomial<F> =
+            limbs_from_vec::<F, N>(P::to_limbs_field::<F>(modulus)).into();
+        let p_result: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(result)).into();
+        let p_carry: Polynomial<F> = limbs_from_vec::<F, N>(P::to_limbs_field::<F>(carry)).into();
 
         // Compute the vanishing polynomial.
         let p_inner_product = p_a_vec
@@ -75,13 +78,13 @@ impl<F: PrimeField32> FieldInnerProductCols<F> {
     }
 }
 
-impl<V: Copy> FieldInnerProductCols<V> {
+impl<V: Copy, const N: usize, const M: usize> FieldInnerProductCols<V, N, M> {
     #[allow(unused_variables)]
     pub fn eval<AB: SP1AirBuilder<Var = V>, P: FieldParameters>(
         &self,
         builder: &mut AB,
-        a: &[Limbs<AB::Var>],
-        b: &[Limbs<AB::Var>],
+        a: &[Limbs<AB::Var, N>],
+        b: &[Limbs<AB::Var, N>],
     ) where
         V: Into<AB::Expr>,
     {
@@ -108,7 +111,7 @@ impl<V: Copy> FieldInnerProductCols<V> {
         let p_witness_low = self.witness_low.iter().into();
         let p_witness_high = self.witness_high.iter().into();
 
-        eval_field_operation::<AB, P>(builder, &p_vanishing, &p_witness_low, &p_witness_high);
+        eval_field_operation::<AB, N, P>(builder, &p_vanishing, &p_witness_low, &p_witness_high);
     }
 }
 
@@ -123,7 +126,7 @@ mod tests {
     use crate::air::MachineAir;
 
     use crate::utils::ec::edwards::ed25519::Ed25519BaseField;
-    use crate::utils::ec::field::FieldParameters;
+    use crate::utils::ec::field::{limbs_from_vec, FieldParameters};
     use crate::utils::{pad_to_power_of_two, BabyBearPoseidon2, StarkUtils};
     use crate::utils::{uni_stark_prove as prove, uni_stark_verify as verify};
     use crate::{air::SP1AirBuilder, runtime::ExecutionRecord};
@@ -138,13 +141,16 @@ mod tests {
     use sp1_derive::AlignedBorrow;
 
     #[derive(AlignedBorrow, Debug, Clone)]
-    pub struct TestCols<T> {
-        pub a: [Limbs<T>; 1],
-        pub b: [Limbs<T>; 1],
-        pub a_ip_b: FieldInnerProductCols<T>,
+    pub struct TestCols<T, const N: usize, const M: usize> {
+        pub a: [Limbs<T, N>; 1],
+        pub b: [Limbs<T, N>; 1],
+        pub a_ip_b: FieldInnerProductCols<T, N, M>,
     }
 
-    pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8>>();
+    const NUM_LIMBS: usize = 32;
+    const NUM_WITNESS_LIMBS: usize = 2 * NUM_LIMBS - 2;
+
+    pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8, NUM_LIMBS, NUM_WITNESS_LIMBS>>();
 
     struct FieldIpChip<P: FieldParameters> {
         pub _phantom: std::marker::PhantomData<P>,
@@ -188,9 +194,10 @@ mod tests {
                 .iter()
                 .map(|(a, b)| {
                     let mut row = [F::zero(); NUM_TEST_COLS];
-                    let cols: &mut TestCols<F> = row.as_mut_slice().borrow_mut();
-                    cols.a[0] = P::to_limbs_field::<F>(&a[0]);
-                    cols.b[0] = P::to_limbs_field::<F>(&b[0]);
+                    let cols: &mut TestCols<F, NUM_LIMBS, NUM_WITNESS_LIMBS> =
+                        row.as_mut_slice().borrow_mut();
+                    cols.a[0] = limbs_from_vec::<F, NUM_LIMBS>(P::to_limbs_field::<F>(&a[0]));
+                    cols.b[0] = limbs_from_vec::<F, NUM_LIMBS>(P::to_limbs_field::<F>(&b[0]));
                     cols.a_ip_b.populate::<P>(a, b);
                     row
                 })
@@ -220,7 +227,8 @@ mod tests {
     {
         fn eval(&self, builder: &mut AB) {
             let main = builder.main();
-            let local: &TestCols<AB::Var> = main.row_slice(0).borrow();
+            let local: &TestCols<AB::Var, NUM_LIMBS, NUM_WITNESS_LIMBS> =
+                main.row_slice(0).borrow();
             local.a_ip_b.eval::<AB, P>(builder, &local.a, &local.b);
 
             // A dummy constraint to keep the degree 3.