refactor: generic size limbs for curve precompiles

core/Cargo.toml

@@ -39,6 +39,8 @@ anyhow = "1.0.79"
 blake3 = "1.5"
 blake3-zkvm = { git = "https://github.com/sp1-patches/BLAKE3.git", branch = "patch-blake3_zkvm/v.1.0.0" }
 cfg-if = "1.0.0"
+generic-array = { version = "1.0.0", features = ["alloc"] }
+typenum = "1.17.0"
 clap = { version = "4.4.0", features = ["derive"] }
 curve25519-dalek = { version = "=4.0.0" }
 elliptic-curve = "0.13.8"

core/src/lib.rs

@@ -8,8 +8,10 @@
     clippy::type_complexity,
     clippy::unnecessary_unwrap,
     clippy::default_constructed_unit_structs,
-    clippy::box_default
+    clippy::box_default,
+    incomplete_features
 )]
+#![feature(generic_const_exprs)]
 
 extern crate alloc;
 

core/src/operations/field/field_den.rs

@@ -1,10 +1,10 @@
-use super::params::Limbs;
-use super::params::NUM_WITNESS_LIMBS;
+use super::params::{Limbs, NumLimbs};
 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 crate::utils::ec::field::{limbs_from_vec, FieldParameters};
+
 use num::BigUint;
 use p3_field::PrimeField32;
 use sp1_derive::AlignedBorrow;
@@ -19,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, N: NumLimbs> {
     /// 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::Limbs>,
+    pub(crate) carry: Limbs<T, N::Limbs>,
+    pub(crate) witness_low: Limbs<T, N::Witness>,
+    pub(crate) witness_high: Limbs<T, N::Witness>,
 }
 
-impl<F: PrimeField32> FieldDenCols<F> {
+impl<F: PrimeField32, N: NumLimbs> FieldDenCols<F, N> {
     pub fn populate<P: FieldParameters>(
         &mut self,
         a: &BigUint,
@@ -53,11 +53,13 @@ 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::Limbs>(P::to_limbs_field::<F>(a)).into();
+        let p_b: Polynomial<F> = limbs_from_vec::<F, N::Limbs>(P::to_limbs_field::<F>(b)).into();
+        let p_p: Polynomial<F> = limbs_from_vec::<F, N::Limbs>(P::to_limbs_field::<F>(&p)).into();
+        let p_result: Polynomial<F> =
+            limbs_from_vec::<F, N::Limbs>(P::to_limbs_field::<F>(&result)).into();
+        let p_carry: Polynomial<F> =
+            limbs_from_vec::<F, N::Limbs>(P::to_limbs_field::<F>(&carry)).into();
 
         // Compute the vanishing polynomial.
         let vanishing_poly = if sign {
@@ -76,20 +78,23 @@ impl<F: PrimeField32> FieldDenCols<F> {
 
         self.result = p_result.into();
         self.carry = p_carry.into();
-        self.witness_low = p_witness_low.try_into().unwrap();
-        self.witness_high = p_witness_high.try_into().unwrap();
+        self.witness_low = Limbs(p_witness_low.try_into().unwrap());
+        self.witness_high = Limbs(p_witness_high.try_into().unwrap());
 
         result
     }
 }
 
-impl<V: Copy> FieldDenCols<V> {
+impl<V: Copy, N: NumLimbs> FieldDenCols<V, N>
+where
+    Limbs<V, N::Limbs>: Copy,
+{
     #[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::Limbs>,
+        b: &Limbs<AB::Var, N::Limbs>,
         sign: bool,
     ) where
         V: Into<AB::Expr>,
@@ -115,8 +120,8 @@ impl<V: Copy> FieldDenCols<V> {
 
         let p_vanishing = p_lhs_minus_rhs - &p_carry * &p_limbs;
 
-        let p_witness_low = self.witness_low.iter().into();
-        let p_witness_high = self.witness_high.iter().into();
+        let p_witness_low = self.witness_low.0.iter().into();
+        let p_witness_high = self.witness_high.0.iter().into();
 
         eval_field_operation::<AB, P>(builder, &p_vanishing, &p_witness_low, &p_witness_high);
     }
@@ -127,14 +132,16 @@ mod tests {
     use num::BigUint;
     use p3_air::BaseAir;
     use p3_field::{Field, PrimeField32};
+    use typenum::U32;
 
     use super::{FieldDenCols, Limbs};
 
     use crate::air::MachineAir;
 
+    use crate::operations::field::params::NumLimbs32;
     use crate::stark::StarkGenericConfig;
     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::BabyBearPoseidon2;
     use crate::utils::{uni_stark_prove as prove, uni_stark_verify as verify};
     use crate::{air::SP1AirBuilder, runtime::ExecutionRecord};
@@ -147,11 +154,14 @@ mod tests {
     use p3_matrix::MatrixRowSlices;
     use rand::thread_rng;
     use sp1_derive::AlignedBorrow;
-    #[derive(AlignedBorrow, Debug, Clone)]
+
+    type Limbs32 = U32;
+
+    #[derive(Debug, Clone, AlignedBorrow)]
     pub struct TestCols<T> {
-        pub a: Limbs<T>,
-        pub b: Limbs<T>,
-        pub a_den_b: FieldDenCols<T>,
+        pub a: Limbs<T, Limbs32>,
+        pub b: Limbs<T, Limbs32>,
+        pub a_den_b: FieldDenCols<T, NumLimbs32>,
     }
 
     pub const NUM_TEST_COLS: usize = size_of::<TestCols<u8>>();
@@ -207,8 +217,8 @@ mod tests {
                 .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);
+                    cols.a = limbs_from_vec::<F, Limbs32>(P::to_limbs_field::<F>(a));
+                    cols.b = limbs_from_vec::<F, Limbs32>(P::to_limbs_field::<F>(b));
                     cols.a_den_b.populate::<P>(a, b, self.sign);
                     row
                 })
@@ -262,7 +272,7 @@ mod tests {
     }
 
     #[test]
-    fn prove_babybear() {
+    fn prove_field() {
         let config = BabyBearPoseidon2::new();
         let mut challenger = config.challenger();