chore: make FpOperations use PrimeField32 to clean up type inconsistency hacks

core/src/operations/field/fp_den.rs

@@ -8,8 +8,7 @@ use crate::utils::ec::field::FieldParameters;
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
 use num::BigUint;
-use p3_baby_bear::BabyBear;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use std::fmt::Debug;
 use valida_derive::AlignedBorrow;
 
@@ -28,18 +27,13 @@ pub struct FpDenCols<T> {
     pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
 }
 
-impl<F: Field> FpDenCols<F> {
+impl<F: PrimeField32> FpDenCols<F> {
     pub fn populate<P: FieldParameters>(
         &mut self,
         a: &BigUint,
         b: &BigUint,
         sign: bool,
     ) -> BigUint {
-        /// TODO: This operation relies on `F` being a PrimeField32, but our traits do not
-        /// support that. This is a hack, since we always use BabyBear, to get around that, but
-        /// all operations using "PF" should use "F" in the future.
-        type PF = BabyBear;
-
         let p = P::modulus();
         let minus_b_int = &p - b;
         let b_signed = if sign { b.clone() } else { minus_b_int };
@@ -59,11 +53,11 @@ impl<F: Field> FpDenCols<F> {
         debug_assert!(carry < p);
         debug_assert_eq!(&carry * &p, &equation_lhs - &equation_rhs);
 
-        let p_a: Polynomial<PF> = P::to_limbs_field::<PF>(a).into();
-        let p_b: Polynomial<PF> = P::to_limbs_field::<PF>(b).into();
-        let p_p: Polynomial<PF> = P::to_limbs_field::<PF>(&p).into();
-        let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(&result).into();
-        let p_carry: Polynomial<PF> = P::to_limbs_field::<PF>(&carry).into();
+        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();
 
         // Compute the vanishing polynomial.
         let vanishing_poly = if sign {
@@ -132,7 +126,7 @@ impl<V: Copy> FpDenCols<V> {
 mod tests {
     use num::BigUint;
     use p3_air::BaseAir;
-    use p3_field::Field;
+    use p3_field::{Field, PrimeField32};
 
     use super::{FpDenCols, Limbs};
     use crate::air::MachineAir;
@@ -173,7 +167,7 @@ mod tests {
         }
     }
 
-    impl<F: Field, P: FieldParameters> MachineAir<F> for FpDenChip<P> {
+    impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for FpDenChip<P> {
         fn name(&self) -> String {
             "FpDen".to_string()
         }

core/src/operations/field/fp_inner_product.rs

@@ -9,9 +9,7 @@ use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
 use num::BigUint;
 use num::Zero;
-use p3_baby_bear::BabyBear;
-use p3_field::AbstractField;
-use p3_field::Field;
+use p3_field::{AbstractField, PrimeField32};
 use std::fmt::Debug;
 use valida_derive::AlignedBorrow;
 
@@ -28,21 +26,12 @@ pub struct FpInnerProductCols<T> {
     pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
 }
 
-impl<F: Field> FpInnerProductCols<F> {
+impl<F: PrimeField32> FpInnerProductCols<F> {
     pub fn populate<P: FieldParameters>(&mut self, a: &[BigUint], b: &[BigUint]) -> BigUint {
-        /// TODO: This operation relies on `F` being a PrimeField32, but our traits do not
-        /// support that. This is a hack, since we always use BabyBear, to get around that, but
-        /// all operations using "PF" should use "F" in the future.
-        type PF = BabyBear;
-
-        let p_a_vec: Vec<Polynomial<PF>> = a
-            .iter()
-            .map(|x| P::to_limbs_field::<PF>(x).into())
-            .collect();
-        let p_b_vec: Vec<Polynomial<PF>> = b
-            .iter()
-            .map(|x| P::to_limbs_field::<PF>(x).into())
-            .collect();
+        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 modulus = &P::modulus();
         let inner_product = a
@@ -56,15 +45,15 @@ impl<F: Field> FpInnerProductCols<F> {
         assert!(carry < &(2u32 * modulus));
         assert_eq!(carry * modulus, inner_product - result);
 
-        let p_modulus: Polynomial<PF> = P::to_limbs_field::<PF>(modulus).into();
-        let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(result).into();
-        let p_carry: Polynomial<PF> = P::to_limbs_field::<PF>(carry).into();
+        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();
 
         // Compute the vanishing polynomial.
         let p_inner_product = p_a_vec
             .into_iter()
             .zip(p_b_vec)
-            .fold(Polynomial::<PF>::new(vec![PF::zero()]), |acc, (c, d)| {
+            .fold(Polynomial::<F>::new(vec![F::zero()]), |acc, (c, d)| {
                 acc + &c * &d
             });
         let p_vanishing = p_inner_product - &p_result - &p_carry * &p_modulus;
@@ -127,7 +116,7 @@ impl<V: Copy> FpInnerProductCols<V> {
 mod tests {
     use num::BigUint;
     use p3_air::BaseAir;
-    use p3_field::Field;
+    use p3_field::{Field, PrimeField32};
 
     use super::{FpInnerProductCols, Limbs};
     use crate::air::MachineAir;
@@ -167,7 +156,7 @@ mod tests {
         }
     }
 
-    impl<F: Field, P: FieldParameters> MachineAir<F> for FpIpChip<P> {
+    impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for FpIpChip<P> {
         fn name(&self) -> String {
             "FpInnerProduct".to_string()
         }

core/src/operations/field/fp_op.rs

@@ -9,8 +9,7 @@ use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
 use num::{BigUint, Zero};
 use p3_air::AirBuilder;
-use p3_baby_bear::BabyBear;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use std::fmt::Debug;
 use valida_derive::AlignedBorrow;
 
@@ -35,18 +34,13 @@ pub struct FpOpCols<T> {
     pub(crate) witness_high: [T; NUM_WITNESS_LIMBS],
 }
 
-impl<F: Field> FpOpCols<F> {
+impl<F: PrimeField32> FpOpCols<F> {
     pub fn populate<P: FieldParameters>(
         &mut self,
         a: &BigUint,
         b: &BigUint,
         op: FpOperation,
     ) -> BigUint {
-        /// TODO: This operation relies on `F` being a PrimeField32, but our traits do not
-        /// support that. This is a hack, since we always use BabyBear, to get around that, but
-        /// all operations using "PF" should use "F" in the future.
-        type PF = BabyBear;
-
         if b == &BigUint::zero() && op == FpOperation::Div {
             // Division by 0 is allowed only when dividing 0 so that padded rows can be all 0.
             assert_eq!(
@@ -67,7 +61,7 @@ impl<F: Field> FpOpCols<F> {
             // Note that this reversal means we have to flip result, a correspondingly in
             // the `eval` function.
             self.populate::<P>(&result, b, FpOperation::Add);
-            let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(&result).into();
+            let p_result: Polynomial<F> = P::to_limbs_field::<F>(&result).into();
             self.result = convert_polynomial(p_result);
             return result;
         }
@@ -85,13 +79,13 @@ impl<F: Field> FpOpCols<F> {
             // Note that this reversal means we have to flip result, a correspondingly in the `eval`
             // function.
             self.populate::<P>(&result, b, FpOperation::Mul);
-            let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(&result).into();
+            let p_result: Polynomial<F> = P::to_limbs_field::<F>(&result).into();
             self.result = convert_polynomial(p_result);
             return result;
         }
 
-        let p_a: Polynomial<PF> = P::to_limbs_field::<PF>(a).into();
-        let p_b: Polynomial<PF> = P::to_limbs_field::<PF>(b).into();
+        let p_a: Polynomial<F> = P::to_limbs_field::<F>(a).into();
+        let p_b: Polynomial<F> = P::to_limbs_field::<F>(b).into();
 
         // Compute field addition in the integers.
         let modulus = &P::modulus();
@@ -109,17 +103,17 @@ impl<F: Field> FpOpCols<F> {
         }
 
         // Make little endian polynomial limbs.
-        let p_modulus: Polynomial<PF> = P::to_limbs_field::<PF>(modulus).into();
-        let p_result: Polynomial<PF> = P::to_limbs_field::<PF>(&result).into();
-        let p_carry: Polynomial<PF> = P::to_limbs_field::<PF>(&carry).into();
+        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();
 
         // Compute the vanishing polynomial.
         let p_op = match op {
             FpOperation::Add => &p_a + &p_b,
             FpOperation::Mul => &p_a * &p_b,
             FpOperation::Sub | FpOperation::Div => unreachable!(),
         };
-        let p_vanishing: Polynomial<PF> = &p_op - &p_result - &p_carry * &p_modulus;
+        let p_vanishing: Polynomial<F> = &p_op - &p_result - &p_carry * &p_modulus;
         debug_assert_eq!(p_vanishing.degree(), P::NB_WITNESS_LIMBS);
 
         let p_witness = compute_root_quotient_and_shift(
@@ -179,7 +173,7 @@ impl<V: Copy> FpOpCols<V> {
 mod tests {
     use num::BigUint;
     use p3_air::BaseAir;
-    use p3_field::Field;
+    use p3_field::{Field, PrimeField32};
 
     use super::{FpOpCols, FpOperation, Limbs};
     use crate::air::MachineAir;
@@ -221,7 +215,7 @@ mod tests {
         }
     }
 
-    impl<F: Field, P: FieldParameters> MachineAir<F> for FpOpChip<P> {
+    impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for FpOpChip<P> {
         fn name(&self) -> String {
             format!("FpOp{:?}", self.operation)
         }

core/src/operations/field/fp_sqrt.rs

@@ -5,7 +5,7 @@ use crate::utils::ec::field::FieldParameters;
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
 use num::BigUint;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use std::fmt::Debug;
 use valida_derive::AlignedBorrow;
 
@@ -21,7 +21,7 @@ pub struct FpSqrtCols<T> {
     pub multiplication: FpOpCols<T>,
 }
 
-impl<F: Field> FpSqrtCols<F> {
+impl<F: PrimeField32> FpSqrtCols<F> {
     /// Populates the trace.
     ///
     /// `P` is the parameter of the field that each limb lives in.
@@ -79,7 +79,7 @@ impl<V: Copy> FpSqrtCols<V> {
 mod tests {
     use num::{BigUint, One, Zero};
     use p3_air::BaseAir;
-    use p3_field::Field;
+    use p3_field::{Field, PrimeField32};
 
     use super::{FpSqrtCols, Limbs};
     use crate::air::MachineAir;
@@ -117,7 +117,7 @@ mod tests {
         }
     }
 
-    impl<F: Field, P: FieldParameters> MachineAir<F> for EdSqrtChip<P> {
+    impl<F: PrimeField32, P: FieldParameters> MachineAir<F> for EdSqrtChip<P> {
         fn name(&self) -> String {
             "EdSqrtChip".to_string()
         }

core/src/operations/field/params.rs

@@ -1,6 +1,4 @@
 use crate::air::Polynomial;
-use p3_baby_bear::BabyBear;
-use p3_field::Field;
 use p3_field::PrimeField32;
 use std::fmt::Debug;
 use std::ops::Index;
@@ -58,7 +56,7 @@ impl<'a, T: Debug + Default + Clone> From<Iter<'a, T>> for Limbs<T> {
 }
 
 // TODO: we probably won't need this in the future when we do things properly.
-pub fn convert_polynomial<F: Field>(value: Polynomial<BabyBear>) -> Limbs<F> {
+pub fn convert_polynomial<F: PrimeField32>(value: Polynomial<F>) -> Limbs<F> {
     let inner_u8 = value
         .as_coefficients()
         .iter()
@@ -69,8 +67,7 @@ pub fn convert_polynomial<F: Field>(value: Polynomial<BabyBear>) -> Limbs<F> {
     Limbs(inner)
 }
 
-// TODO: we probably won't need this in the future when we do things properly.
-pub fn convert_vec<F: Field>(value: Vec<BabyBear>) -> Vec<F> {
+pub fn convert_vec<F: PrimeField32>(value: Vec<F>) -> Vec<F> {
     value
         .iter()
         .map(|x| x.as_canonical_u32() as u8)

core/src/syscall/precompiles/edwards/ed_add.rs

@@ -27,7 +27,7 @@ use num::Zero;
 use p3_air::AirBuilder;
 use p3_air::{Air, BaseAir};
 use p3_field::AbstractField;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_matrix::MatrixRowSlices;
 use p3_maybe_rayon::prelude::IntoParallelRefIterator;
@@ -74,7 +74,7 @@ impl<E: EllipticCurve + EdwardsParameters> EdAddAssignChip<E> {
             _marker: PhantomData,
         }
     }
-    fn populate_fp_ops<F: Field>(
+    fn populate_fp_ops<F: PrimeField32>(
         cols: &mut EdAddAssignCols<F>,
         p_x: BigUint,
         p_y: BigUint,
@@ -121,7 +121,7 @@ impl<E: EllipticCurve + EdwardsParameters> Syscall for EdAddAssignChip<E> {
     }
 }
 
-impl<F: Field, E: EllipticCurve + EdwardsParameters> MachineAir<F> for EdAddAssignChip<E> {
+impl<F: PrimeField32, E: EllipticCurve + EdwardsParameters> MachineAir<F> for EdAddAssignChip<E> {
     fn name(&self) -> String {
         "EdAddAssign".to_string()
     }

core/src/syscall/precompiles/edwards/ed_decompress.rs

@@ -32,7 +32,7 @@ use num::One;
 use num::Zero;
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use p3_matrix::MatrixRowSlices;
 use std::marker::PhantomData;
 
@@ -77,7 +77,7 @@ pub struct EdDecompressCols<T> {
     pub(crate) neg_x: FpOpCols<T>,
 }
 
-impl<F: Field> EdDecompressCols<F> {
+impl<F: PrimeField32> EdDecompressCols<F> {
     pub fn populate<P: FieldParameters, E: EdwardsParameters>(
         &mut self,
         event: EdDecompressEvent,
@@ -267,7 +267,7 @@ impl<E: EdwardsParameters> EdDecompressChip<E> {
     }
 }
 
-impl<F: Field, E: EdwardsParameters> MachineAir<F> for EdDecompressChip<E> {
+impl<F: PrimeField32, E: EdwardsParameters> MachineAir<F> for EdDecompressChip<E> {
     fn name(&self) -> String {
         "EdDecompress".to_string()
     }

core/src/syscall/precompiles/k256/decompress.rs

@@ -35,7 +35,7 @@ use num::Zero;
 use p3_air::AirBuilder;
 use p3_air::{Air, BaseAir};
 use p3_field::AbstractField;
-use p3_field::Field;
+use p3_field::PrimeField32;
 use p3_matrix::MatrixRowSlices;
 use std::str::FromStr;
 
@@ -152,7 +152,7 @@ pub struct K256DecompressCols<T> {
     pub(crate) y_least_bits: [T; 8],
 }
 
-impl<F: Field> K256DecompressCols<F> {
+impl<F: PrimeField32> K256DecompressCols<F> {
     pub fn populate(&mut self, event: K256DecompressEvent, shard: &mut ExecutionRecord) {
         let mut new_field_events = Vec::new();
         self.is_real = F::from_bool(true);
@@ -289,7 +289,7 @@ impl<V: Copy> K256DecompressCols<V> {
     }
 }
 
-impl<F: Field> MachineAir<F> for K256DecompressChip {
+impl<F: PrimeField32> MachineAir<F> for K256DecompressChip {
     fn name(&self) -> String {
         "K256Decompress".to_string()
     }