Pin prime subfield koalabear

crates/backend/air/Cargo.toml

@@ -5,4 +5,5 @@ edition.workspace = true
 
 [dependencies]
 field = { path = "../field", package = "field" }
+koala-bear = { path = "../koala-bear", package = "koala-bear" }
 poly = { path = "../poly", package = "poly" }

crates/backend/air/src/constraint_folder/normal.rs

@@ -3,7 +3,7 @@ use field::*;
 use poly::*;
 
 #[derive(Debug)]
-pub struct ConstraintFolder<'a, IF, EF: ExtensionField<PF<EF>>, ExtraData: AlphaPowers<EF>> {
+pub struct ConstraintFolder<'a, IF, EF: KoalaBearExtension, ExtraData: AlphaPowers<EF>> {
     pub flat: &'a [IF],
     pub shift: &'a [IF],
     pub extra_data: &'a ExtraData,
@@ -13,7 +13,7 @@ pub struct ConstraintFolder<'a, IF, EF: ExtensionField<PF<EF>>, ExtraData: Alpha
 
 impl<'a, IF, EF, ExtraData> ConstraintFolder<'a, IF, EF, ExtraData>
 where
-    EF: ExtensionField<PF<EF>>,
+    EF: KoalaBearExtension,
     ExtraData: AlphaPowers<EF>,
 {
     pub fn new(flat: &'a [IF], shift: &'a [IF], extra_data: &'a ExtraData) -> Self {
@@ -30,7 +30,7 @@ where
 impl<'a, IF, EF, ExtraData> AirBuilder for ConstraintFolder<'a, IF, EF, ExtraData>
 where
     IF: Algebra<PF<EF>> + 'static,
-    EF: Field + ExtensionField<PF<EF>> + Mul<IF, Output = EF> + Add<IF, Output = EF>,
+    EF: Field + KoalaBearExtension + Mul<IF, Output = EF> + Add<IF, Output = EF>,
     ExtraData: AlphaPowers<EF>,
 {
     type F = PF<EF>;

crates/backend/air/src/constraint_folder/packed.rs

@@ -1,5 +1,6 @@
 use crate::*;
 use field::*;
+use koala_bear::KoalaBear;
 use poly::*;
 
 #[derive(Debug)]
@@ -38,7 +39,7 @@ where
 
 impl<'a, IF, EF, ExtraData> AirBuilder for ConstraintFolderPacked<'a, IF, EF, ExtraData>
 where
-    IF: Algebra<PFPacking<EF>> + 'static,
+    IF: Algebra<PFPacking<EF>> + Algebra<KoalaBear> + 'static,
     EF: Field + ExtensionField<PF<EF>>,
     EFPacking<EF>: PrimeCharacteristicRing + Mul<IF, Output = EFPacking<EF>> + Add<IF, Output = EFPacking<EF>>,
     ExtraData: AlphaPowers<EF>,

crates/backend/air/src/lib.rs

@@ -2,6 +2,7 @@
 
 use core::ops::{Add, Mul, Sub};
 use field::{Algebra, PrimeCharacteristicRing};
+use koala_bear::KoalaBear;
 
 mod symbolic;
 pub use symbolic::*;
@@ -36,7 +37,7 @@ pub trait AirBuilder: Sized {
     type F: PrimeCharacteristicRing + 'static;
     /// Intermediate field: equals F in base-field rounds, EF in extension rounds
     /// (or their respective SIMD packings).
-    type IF: Algebra<Self::F> + 'static;
+    type IF: Algebra<Self::F> + Algebra<KoalaBear> + 'static;
     /// Always the extension field (or its SIMD packing).
     type EF: PrimeCharacteristicRing
         + 'static

crates/backend/air/src/symbolic.rs

@@ -7,6 +7,7 @@ use core::marker::PhantomData;
 use core::ops::{Add, AddAssign, Deref, Mul, MulAssign, Neg, Sub, SubAssign};
 
 use field::{Algebra, Field, InjectiveMonomial, PrimeCharacteristicRing};
+use koala_bear::KoalaBear;
 
 use crate::{Air, AirBuilder};
 
@@ -284,7 +285,10 @@ impl<F: Field> SymbolicAirBuilder<F> {
     }
 }
 
-impl<F: Field> AirBuilder for SymbolicAirBuilder<F> {
+impl<F: Field> AirBuilder for SymbolicAirBuilder<F>
+where
+    SymbolicExpression<F>: Algebra<KoalaBear>,
+{
     type F = F;
     type IF = SymbolicExpression<F>;
     type EF = SymbolicExpression<F>;
@@ -325,6 +329,7 @@ pub type SymbolicAirData<F> = (
 pub fn get_symbolic_constraints_and_bus_data_values<F: Field, A: Air>(air: &A) -> SymbolicAirData<F>
 where
     A::ExtraData: Default,
+    SymbolicExpression<F>: Algebra<KoalaBear>,
 {
     let mut builder = SymbolicAirBuilder::<F>::new(air.n_columns(), air.n_shift_columns());
     air.eval(&mut builder, &Default::default());

crates/backend/fiat-shamir/src/prover.rs

@@ -3,8 +3,9 @@ use crate::{MerklePaths, PrunedMerklePaths, *};
 use field::Field;
 use field::PackedValue;
 use field::PrimeCharacteristicRing;
+use field::PrimeField64;
 use field::integers::QuotientMap;
-use field::{ExtensionField, PrimeField64};
+use koala_bear::KoalaBearExtension;
 use koala_bear::symmetric::Permutation;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::time::Duration;
@@ -24,16 +25,13 @@ pub fn reset_pow_grinding_time() {
 }
 
 #[derive(Debug)]
-pub struct ProverState<EF: ExtensionField<PF<EF>>, P> {
+pub struct ProverState<EF: KoalaBearExtension, P> {
     challenger: Challenger<PF<EF>, P>,
     transcript: Vec<PF<EF>>,
     merkle_paths: Vec<PrunedMerklePaths<PF<EF>, PF<EF>>>,
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]>> ProverState<EF, P>
-where
-    PF<EF>: PrimeField64,
-{
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]>> ProverState<EF, P> {
     #[must_use]
     pub fn new(permutation: P, capacity: [PF<EF>; CAPACITY]) -> Self {
         assert!(EF::DIMENSION <= RATE);
@@ -52,10 +50,7 @@ where
     }
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]>> ChallengeSampler<EF> for ProverState<EF, P>
-where
-    PF<EF>: PrimeField64,
-{
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]>> ChallengeSampler<EF> for ProverState<EF, P> {
     fn sample_vec(&mut self, len: usize) -> Vec<EF> {
         sample_vec(&mut self.challenger, len)
     }
@@ -65,10 +60,8 @@ where
     }
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]> + Permutation<[<PF<EF> as Field>::Packing; WIDTH]>>
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]> + Permutation<[<PF<EF> as Field>::Packing; WIDTH]>>
     FSProver<EF> for ProverState<EF, P>
-where
-    PF<EF>: PrimeField64,
 {
     fn add_base_scalars(&mut self, scalars: &[PF<EF>]) {
         self.challenger.observe_many(scalars);

crates/backend/fiat-shamir/src/traits.rs

@@ -1,4 +1,4 @@
-use field::ExtensionField;
+use koala_bear::KoalaBearExtension;
 
 use crate::{
     MerkleOpening, MerklePath, PF, ProofError, ProofResult, flatten_scalars_to_base, pack_scalars_to_extension,
@@ -12,7 +12,7 @@ pub trait ChallengeSampler<EF> {
     fn sample_in_range(&mut self, bits: usize, n_samples: usize) -> Vec<usize>;
 }
 
-pub trait FSProver<EF: ExtensionField<PF<EF>>>: ChallengeSampler<EF> {
+pub trait FSProver<EF: KoalaBearExtension>: ChallengeSampler<EF> {
     fn state(&self) -> String;
     fn add_base_scalars(&mut self, scalars: &[PF<EF>]);
     fn observe_scalars(&mut self, scalars: &[PF<EF>]);
@@ -43,7 +43,7 @@ pub trait FSProver<EF: ExtensionField<PF<EF>>>: ChallengeSampler<EF> {
     }
 }
 
-pub trait FSVerifier<EF: ExtensionField<PF<EF>>>: ChallengeSampler<EF> {
+pub trait FSVerifier<EF: KoalaBearExtension>: ChallengeSampler<EF> {
     fn state(&self) -> String;
     fn next_base_scalars_vec(&mut self, n: usize) -> Result<Vec<PF<EF>>, ProofError>;
     fn observe_scalars(&mut self, scalars: &[PF<EF>]);

crates/backend/fiat-shamir/src/verifier.rs

@@ -5,17 +5,17 @@ use crate::{
     *,
 };
 use field::PrimeCharacteristicRing;
-use field::{ExtensionField, PrimeField64};
+use field::PrimeField64;
+use koala_bear::KoalaBearExtension;
 use koala_bear::symmetric::Permutation;
 use koala_bear::{KoalaBear, default_koalabear_poseidon1_16};
-use std::any::TypeId;
 use std::collections::VecDeque;
 use std::iter::repeat_n;
 use symetric::CAPACITY;
 use symetric::RATE;
 use symetric::WIDTH;
 
-pub struct VerifierState<EF: ExtensionField<PF<EF>>, P> {
+pub struct VerifierState<EF: KoalaBearExtension, P> {
     challenger: Challenger<PF<EF>, P>,
     transcript: Vec<PF<EF>>,
     transcript_offset: usize,
@@ -25,10 +25,7 @@ pub struct VerifierState<EF: ExtensionField<PF<EF>>, P> {
     raw_transcript: Vec<PF<EF>>, // reconstructed during the proof verification, it's the format that the zkVM recursion program expects (no Merkle pruning, no sumcheck optimization to send less data, etc)
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]>> VerifierState<EF, P>
-where
-    PF<EF>: PrimeField64,
-{
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]>> VerifierState<EF, P> {
     pub fn new(proof: Proof<PF<EF>>, permutation: P, capacity: [PF<EF>; CAPACITY]) -> Result<Self, ProofError> {
         Ok(Self {
             challenger: Challenger::new(permutation, capacity),
@@ -75,16 +72,12 @@ where
         Ok(scalars)
     }
 
-    #[allow(clippy::missing_transmute_annotations)]
     fn restore_merkle_paths(
-        paths: PrunedMerklePaths<PF<EF>, PF<EF>>,
+        paths: PrunedMerklePaths<KoalaBear, KoalaBear>,
         indices: &[usize],
         merkle_height: usize,
         leaf_len: usize,
-    ) -> Option<Vec<MerkleOpening<PF<EF>>>> {
-        assert_eq!(TypeId::of::<PF<EF>>(), TypeId::of::<KoalaBear>());
-        // SAFETY: We've confirmed PF<EF> == KoalaBear
-        let paths: PrunedMerklePaths<KoalaBear, KoalaBear> = unsafe { std::mem::transmute(paths) };
+    ) -> Option<Vec<MerkleOpening<KoalaBear>>> {
         let perm = default_koalabear_poseidon1_16();
         let hash_fn = |data: &[KoalaBear]| symetric::hash_slice_rtl::<_, _, 16, 8, DIGEST_LEN_FE>(&perm, data);
         let combine_fn = |left: &[KoalaBear; DIGEST_LEN_FE], right: &[KoalaBear; DIGEST_LEN_FE]| {
@@ -100,15 +93,11 @@ where
                 path: path.sibling_hashes,
             })
             .collect();
-        // SAFETY: PF<EF> == KoalaBear
-        Some(unsafe { std::mem::transmute(openings) })
+        Some(openings)
     }
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]>> ChallengeSampler<EF> for VerifierState<EF, P>
-where
-    PF<EF>: PrimeField64,
-{
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]>> ChallengeSampler<EF> for VerifierState<EF, P> {
     fn sample_vec(&mut self, len: usize) -> Vec<EF> {
         sample_vec(&mut self.challenger, len)
     }
@@ -117,10 +106,7 @@ where
     }
 }
 
-impl<EF: ExtensionField<PF<EF>>, P: Permutation<[PF<EF>; WIDTH]>> FSVerifier<EF> for VerifierState<EF, P>
-where
-    PF<EF>: PrimeField64,
-{
+impl<EF: KoalaBearExtension, P: Permutation<[PF<EF>; WIDTH]>> FSVerifier<EF> for VerifierState<EF, P> {
     fn state(&self) -> String {
         format!(
             "state {} (offset: {}, merkle_idx: {})",

crates/backend/field/src/field.rs

@@ -6,11 +6,10 @@ use core::fmt::{Debug, Display};
 use core::hash::Hash;
 use core::iter::{Product, Sum};
 use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
-use core::{array, slice};
+use core::slice;
 
 use serde::Serialize;
 use serde::de::DeserializeOwned;
-use utils::iter_array_chunks_padded;
 
 use crate::exponentiation::bits_u64;
 use crate::integers::{QuotientMap, from_integer_types};
@@ -623,112 +622,12 @@ pub trait Algebra<F>:
 // Every ring is an algebra over itself.
 impl<R: PrimeCharacteristicRing> Algebra<R> for R {}
 
-/// A collection of methods designed to help hash field elements.
-///
-/// Most fields will want to reimplement many/all of these methods as the default implementations
-/// are slow and involve converting to/from byte representations.
-pub trait RawDataSerializable: Sized {
-    /// The number of bytes which this field element occupies in memory.
-    /// Must be equal to the length of self.into_bytes().
-    const NUM_BYTES: usize;
-
-    /// Convert a field element into a collection of bytes.
-    #[must_use]
-    fn into_bytes(self) -> impl IntoIterator<Item = u8>;
-
-    /// Convert an iterator of field elements into an iterator of bytes.
-    #[must_use]
-    fn into_byte_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u8> {
-        input.into_iter().flat_map(|elem| elem.into_bytes())
-    }
-
-    /// Convert an iterator of field elements into an iterator of u32s.
-    ///
-    /// If `NUM_BYTES` does not divide `4`, multiple `F`s may be packed together to make a single `u32`. Furthermore,
-    /// if `NUM_BYTES * input.len()` does not divide `4`, the final `u32` will involve padding bytes which are set to `0`.
-    #[must_use]
-    fn into_u32_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u32> {
-        let bytes = Self::into_byte_stream(input);
-        iter_array_chunks_padded(bytes, 0).map(u32::from_le_bytes)
-    }
-
-    /// Convert an iterator of field elements into an iterator of u64s.
-    ///
-    /// If `NUM_BYTES` does not divide `8`, multiple `F`s may be packed together to make a single `u64`. Furthermore,
-    /// if `NUM_BYTES * input.len()` does not divide `8`, the final `u64` will involve padding bytes which are set to `0`.
-    #[must_use]
-    fn into_u64_stream(input: impl IntoIterator<Item = Self>) -> impl IntoIterator<Item = u64> {
-        let bytes = Self::into_byte_stream(input);
-        iter_array_chunks_padded(bytes, 0).map(u64::from_le_bytes)
-    }
-
-    /// Convert an iterator of field element arrays into an iterator of byte arrays.
-    ///
-    /// Converts an element `[F; N]` into the byte array `[[u8; N]; NUM_BYTES]`. This is
-    /// intended for use with vectorized hash functions which use vector operations
-    /// to compute several hashes in parallel.
-    #[must_use]
-    fn into_parallel_byte_streams<const N: usize>(
-        input: impl IntoIterator<Item = [Self; N]>,
-    ) -> impl IntoIterator<Item = [u8; N]> {
-        input.into_iter().flat_map(|vector| {
-            let bytes = vector.map(|elem| elem.into_bytes().into_iter().collect::<Vec<_>>());
-            (0..Self::NUM_BYTES).map(move |i| array::from_fn(|j| bytes[j][i]))
-        })
-    }
-
-    /// Convert an iterator of field element arrays into an iterator of u32 arrays.
-    ///
-    /// Converts an element `[F; N]` into the u32 array `[[u32; N]; NUM_BYTES/4]`. This is
-    /// intended for use with vectorized hash functions which use vector operations
-    /// to compute several hashes in parallel.
-    ///
-    /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose
-    /// operation to get `[Iterator<F>; N]`, calling `into_u32_stream` on each element to get `[Iterator<u32>; N]` and then
-    /// performing another transpose operation to get `Iterator<[u32; N]>`.
-    ///
-    /// If `NUM_BYTES` does not divide `4`, multiple `[F; N]`s may be packed together to make a single `[u32; N]`. Furthermore,
-    /// if `NUM_BYTES * input.len()` does not divide `4`, the final `[u32; N]` will involve padding bytes which are set to `0`.
-    #[must_use]
-    fn into_parallel_u32_streams<const N: usize>(
-        input: impl IntoIterator<Item = [Self; N]>,
-    ) -> impl IntoIterator<Item = [u32; N]> {
-        let bytes = Self::into_parallel_byte_streams(input);
-        iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 4]| {
-            array::from_fn(|i| u32::from_le_bytes(array::from_fn(|j| byte_array[j][i])))
-        })
-    }
-
-    /// Convert an iterator of field element arrays into an iterator of u64 arrays.
-    ///
-    /// Converts an element `[F; N]` into the u64 array `[[u64; N]; NUM_BYTES/8]`. This is
-    /// intended for use with vectorized hash functions which use vector operations
-    /// to compute several hashes in parallel.
-    ///
-    /// This function is guaranteed to be equivalent to starting with `Iterator<[F; N]>` performing a transpose
-    /// operation to get `[Iterator<F>; N]`, calling `into_u64_stream` on each element to get `[Iterator<u64>; N]` and then
-    /// performing another transpose operation to get `Iterator<[u64; N]>`.
-    ///
-    /// If `NUM_BYTES` does not divide `8`, multiple `[F; N]`s may be packed together to make a single `[u64; N]`. Furthermore,
-    /// if `NUM_BYTES * input.len()` does not divide `8`, the final `[u64; N]` will involve padding bytes which are set to `0`.
-    #[must_use]
-    fn into_parallel_u64_streams<const N: usize>(
-        input: impl IntoIterator<Item = [Self; N]>,
-    ) -> impl IntoIterator<Item = [u64; N]> {
-        let bytes = Self::into_parallel_byte_streams(input);
-        iter_array_chunks_padded(bytes, [0; N]).map(|byte_array: [[u8; N]; 8]| {
-            array::from_fn(|i| u64::from_le_bytes(array::from_fn(|j| byte_array[j][i])))
-        })
-    }
-}
-
 /// A field `F`. This permits both modular fields `ℤ/p` along with their field extensions.
 ///
 /// A ring is a field if every element `x` has a unique multiplicative inverse `x^{-1}`
 /// which satisfies `x * x^{-1} = F::ONE`.
 pub trait Field:
     Algebra<Self>
-    + RawDataSerializable
     + Packable
     + 'static
     + Copy