integrate babybear

.cargo/config.toml

@@ -0,0 +1,2 @@
+[build]
+rustflags = ["-C", "target-cpu=native"]

.github/workflows/ci.yml

@@ -16,7 +16,9 @@ jobs:
         with:
           components: rustfmt, clippy
       - run: cargo fmt --all -- --check
-      - run: cargo clippy
+      - run: cargo clippy --all
+        env:
+          RUSTFLAGS: "-Dwarnings -C target-cpu=native -C target-feature=+avx512f"
 
   build-and-test:
     name: Build and Test (${{ matrix.os }}${{ matrix.feature != '' && format(', {0}', matrix.feature) || '' }})
@@ -40,9 +42,7 @@ jobs:
         run: |
           cargo run --bin=dev-setup --release
           cargo build --all-features --release
-          cargo test --all-features --release      
-          cargo test -p arith --release
-          cargo test -p bi-kzg --release
+          cargo test --all-features --release --workspace
 
   benchmark:
     name: Benchmark (${{ matrix.os }}${{ matrix.feature != '' && format(', {0}', matrix.feature) || '' }}, ${{ matrix.field }})
@@ -66,13 +66,13 @@ jobs:
             feature: avx2
             field: gf2ext128
           - os: 7950x3d
-            feature: avx512
+            feature: avx512f
             field: m31ext3
           - os: 7950x3d
-            feature: avx512
+            feature: avx512f
             field: fr
           - os: 7950x3d
-            feature: avx512
+            feature: avx512f
             field: gf2ext128
     steps:
       - uses: actions/checkout@v4
@@ -85,4 +85,3 @@ jobs:
         run: |
           cargo run --bin=dev-setup --release
           cargo run --release -- -t ${{ matrix.os == 'macos-latest' && 2 || 16 }} -f ${{ matrix.field }}
-

.gitignore

@@ -34,6 +34,12 @@ output.txt
 # circuit files
 /data/*
 
+# Generated by SageMath
+*.sage.py
+
 # idea
 .idea/
 /data
+
+# vscode settings
+.vscode/

Cargo.toml

@@ -6,6 +6,11 @@ default-run = "expander-rs" # default
 
 [dependencies]
 arith = { path = "arith" }
+babybear = { path = "arith/babybear" }
+gf2 = { path = "arith/gf2" }
+gf2_128 = { path = "arith/gf2_128" }
+mersenne31 = { path = "arith/mersenne31" }
+
 ark-std.workspace = true
 clap.workspace = true
 env_logger.workspace = true
@@ -35,13 +40,21 @@ name = "dev-setup"
 path = "src/utils.rs"
 
 [features]
+avx256 = [ "arith/avx256" ]
 default = []
 # default = [ "grinding" ]
 grinding = []
-avx256 = ["arith/avx256"]
 
 [workspace]
-members = ["arith", "bi-kzg"]
+members = [
+    "arith", # definitions of Field, Extensions, SIMDS
+    "arith/babybear",
+    "arith/gf2",    
+    "arith/gf2_128",
+    "arith/goldilocks",
+    "arith/mersenne31",
+    "bi-kzg"
+    ]
 
 [workspace.dependencies]
 ark-std = "0.4"
@@ -58,6 +71,10 @@ halo2curves = { git = "https://github.com/PolyhedraZK/halo2curves", default-feat
 ] }
 itertools = "0.13"
 log = "0.4"
+p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", version = "0.1.0", features = ["nightly-features"] }
+p3-field = { git = "https://github.com/Plonky3/Plonky3.git", version = "0.1.0" }
+# This import should be unnecessary but p3-baby-bear fails to turn on "nightly-features" for this dependency
+p3-monty-31 = { git = "https://github.com/Plonky3/Plonky3.git", version = "0.1.0", features = ["nightly-features"] }
 rand = "0.8.5"
 rayon = "1.10"
 sha2 = "0.10.8"
@@ -69,14 +86,24 @@ thiserror = "1.0.63"
 ethnum = "1.5.0"
 
 [[bench]]
-name = "arith"
+name = "bn254_fr"
+harness = false
+path = "arith/benches/fr.rs"
+
+[[bench]]
+name = "gf2"
+harness = false
+path = "arith/gf2/benches/gf2.rs"
+
+[[bench]]
+name = "gf2_128"
 harness = false
-path = "arith/benches/field.rs"
+path = "arith/gf2_128/benches/gf2_128.rs"
 
 [[bench]]
-name = "arith-ext-by-base"
+name = "m31"
 harness = false
-path = "arith/benches/ext_field.rs"
+path = "arith/mersenne31/benches/m31.rs"
 
 [[bench]]
 name = "gkr-hashes"

arith/Cargo.toml

@@ -5,26 +5,25 @@ edition = "2021"
 
 [dependencies]
 ark-std.workspace = true
+criterion.workspace = true
+ethnum.workspace = true
 halo2curves.workspace = true
 log.workspace = true
+p3-baby-bear.workspace = true
+p3-field.workspace = true
+p3-monty-31.workspace = true
 rand.workspace = true
 sha2.workspace = true
 thiserror.workspace = true
-ethnum.workspace = true
+tynm.workspace = true
 
 raw-cpuid = "11.1.0"
 cfg-if = "1.0"
 
 [dev-dependencies]
-tynm.workspace = true
-criterion.workspace = true
-
-[[bench]]
-name = "field"
-harness = false
 
 [[bench]]
-name = "ext_field"
+name = "fr"
 harness = false
 
 [features]

arith/babybear/Cargo.toml

@@ -0,0 +1,14 @@
+[package]
+name = "babybear"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arith = {path = "../"}
+
+ark-std.workspace = true
+ethnum.workspace = true
+p3-baby-bear.workspace = true
+p3-field.workspace = true
+p3-monty-31.workspace = true
+rand.workspace = true

arith/babybear/src/baby_bear.rs

@@ -0,0 +1,208 @@
+use core::{
+    iter::{Product, Sum},
+    ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
+};
+
+use arith::{field_common, Field, FieldForECC, FieldSerde, FieldSerdeResult, SimdField};
+use ark_std::Zero;
+use p3_baby_bear::BabyBear as P3BabyBear;
+use p3_field::{Field as P3Field, PrimeField32};
+use rand::distributions::{Distribution, Standard};
+use std::io::{Read, Write};
+
+mod baby_bearx16;
+pub use baby_bearx16::BabyBearx16;
+
+#[cfg(target_arch = "x86_64")]
+pub(crate) mod baby_bear_avx;
+#[cfg(target_arch = "aarch64")]
+pub(crate) mod baby_bear_neon;
+#[cfg(target_arch = "x86_64")]
+pub use baby_bear_avx::AVXBabyBear;
+
+#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq, PartialOrd, Ord)]
+#[repr(transparent)]
+pub struct BabyBear(P3BabyBear);
+
+pub const BABYBEAR_MODULUS: u32 = 0x78000001;
+
+field_common!(BabyBear);
+
+impl BabyBear {
+    /// Input is provided in canonical form and converted into Montgomery form.
+    pub const fn new(value: u32) -> Self {
+        Self(P3BabyBear::new(value))
+    }
+}
+
+impl FieldSerde for BabyBear {
+    const SERIALIZED_SIZE: usize = 32 / 8;
+
+    #[inline(always)]
+    fn serialize_into<W: Write>(&self, mut writer: W) -> FieldSerdeResult<()> {
+        // Note: BabyBear's impl of as_u32_unchecked() converts to canonical form
+        writer.write_all(self.as_u32_unchecked().to_le_bytes().as_ref())?;
+        Ok(())
+    }
+
+    /// Note: This function performs modular reduction on inputs and
+    /// converts from canonical to Montgomery form.
+    #[inline(always)]
+    fn deserialize_from<R: Read>(mut reader: R) -> FieldSerdeResult<Self> {
+        let mut u = [0u8; Self::SERIALIZED_SIZE];
+        reader.read_exact(&mut u)?;
+        let v = u32::from_le_bytes(u);
+        Ok(Self::from(v))
+    }
+
+    #[inline]
+    fn try_deserialize_from_ecc_format<R: Read>(mut reader: R) -> FieldSerdeResult<Self> {
+        let mut buf = [0u8; 32];
+        reader.read_exact(&mut buf)?;
+        assert!(
+            buf.iter().skip(4).all(|&x| x == 0),
+            "non-zero byte found in witness byte"
+        );
+        Ok(Self::from(u32::from_le_bytes(buf[..4].try_into().unwrap())))
+    }
+}
+
+impl Field for BabyBear {
+    const NAME: &'static str = "Baby Bear Field";
+
+    const SIZE: usize = 32 / 8;
+
+    const FIELD_SIZE: usize = 32;
+
+    const ZERO: Self = BabyBear::new(0);
+
+    const ONE: Self = BabyBear::new(1);
+
+    // See test below
+    const INV_2: Self = BabyBear::new(1006632961);
+
+    #[inline(always)]
+    fn zero() -> Self {
+        Self::ZERO
+    }
+
+    #[inline(always)]
+    fn is_zero(&self) -> bool {
+        *self == Self::ZERO
+    }
+
+    #[inline(always)]
+    fn one() -> Self {
+        Self::ONE
+    }
+
+    /// Uses rejection sampling to avoid bias.
+    fn random_unsafe(mut rng: impl rand::RngCore) -> Self {
+        let dist = Standard;
+        Self(dist.sample(&mut rng))
+    }
+
+    fn random_bool(mut rng: impl rand::RngCore) -> Self {
+        (rng.next_u32() & 1).into()
+    }
+
+    fn exp(&self, exponent: u128) -> Self {
+        let mut e = exponent;
+        let mut res = Self::one();
+        let mut t = *self;
+        while !e.is_zero() {
+            let b = e & 1;
+            if b == 1 {
+                res *= t;
+            }
+            t = t * t;
+            e >>= 1;
+        }
+        res
+    }
+
+    fn inv(&self) -> Option<Self> {
+        self.0.try_inverse().map(Self)
+    }
+
+    /// Converts to canonical form.
+    #[inline(always)]
+    fn as_u32_unchecked(&self) -> u32 {
+        self.0.as_canonical_u32()
+    }
+
+    #[inline(always)]
+    fn from_uniform_bytes(bytes: &[u8; 32]) -> Self {
+        // Note: From<u32> performs modular reduction
+        u32::from_le_bytes(bytes[..4].try_into().unwrap()).into()
+    }
+}
+
+impl FieldForECC for BabyBear {
+    fn modulus() -> ethnum::U256 {
+        ethnum::U256::from(<P3BabyBear as PrimeField32>::ORDER_U32)
+    }
+
+    fn from_u256(x: ethnum::U256) -> Self {
+        Self::new((x % Self::modulus()).as_u32())
+    }
+
+    fn to_u256(&self) -> ethnum::U256 {
+        // Converts to canonical form before casting to U256
+        ethnum::U256::from(self.as_u32_unchecked())
+    }
+}
+
+// TODO: Actual SIMD impl
+// This is a dummy implementation to satisfy trait bounds
+impl SimdField for BabyBear {
+    type Scalar = Self;
+
+    fn scale(&self, challenge: &Self::Scalar) -> Self {
+        self * challenge
+    }
+
+    fn pack(base_vec: &[Self::Scalar]) -> Self {
+        debug_assert!(base_vec.len() == 1);
+        base_vec[0]
+    }
+
+    fn unpack(&self) -> Vec<Self::Scalar> {
+        vec![*self]
+    }
+
+    fn pack_size() -> usize {
+        1
+    }
+}
+
+impl Neg for BabyBear {
+    type Output = Self;
+
+    #[inline(always)]
+    fn neg(self) -> Self::Output {
+        Self(self.0.neg())
+    }
+}
+
+impl From<u32> for BabyBear {
+    #[inline(always)]
+    fn from(value: u32) -> Self {
+        Self::new(value)
+    }
+}
+
+#[inline(always)]
+fn add_internal(a: &BabyBear, b: &BabyBear) -> BabyBear {
+    BabyBear(a.0 + b.0)
+}
+
+#[inline(always)]
+fn sub_internal(a: &BabyBear, b: &BabyBear) -> BabyBear {
+    BabyBear(a.0 - b.0)
+}
+
+#[inline(always)]
+fn mul_internal(a: &BabyBear, b: &BabyBear) -> BabyBear {
+    BabyBear(a.0 * b.0)
+}