From c7fdc6cb2b0dfbfd0e4d6369d86ba0c683c6f78b Mon Sep 17 00:00:00 2001 From: Leo Date: Fri, 16 Aug 2024 15:40:26 +0200 Subject: [PATCH] baby bear field only (#1696) Most minimal attempt at adding BabyBear, only adds the field and field tests, but no powdr tests using the field. --- cli/src/main.rs | 4 +- cli/src/util.rs | 1 + number/Cargo.toml | 14 +- number/src/baby_bear.rs | 467 ++++++++++++++++++++++++++++++++++++++++ number/src/lib.rs | 2 + number/src/traits.rs | 1 + riscv/src/code_gen.rs | 1 + 7 files changed, 485 insertions(+), 5 deletions(-) create mode 100644 number/src/baby_bear.rs diff --git a/cli/src/main.rs b/cli/src/main.rs index 1215844b2..455a870e8 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -8,7 +8,7 @@ use env_logger::{Builder, Target}; use log::LevelFilter; use powdr_backend::BackendType; use powdr_number::{buffered_write_file, read_polys_csv_file, CsvRenderMode}; -use powdr_number::{BigUint, Bn254Field, FieldElement, GoldilocksField}; +use powdr_number::{BabyBearField, BigUint, Bn254Field, FieldElement, GoldilocksField}; use powdr_pipeline::Pipeline; use std::io; use std::path::PathBuf; @@ -56,6 +56,8 @@ fn bind_cli_args( #[derive(Clone, EnumString, EnumVariantNames, Display)] pub enum FieldArgument { + #[strum(serialize = "bb")] + Bb, #[strum(serialize = "gl")] Gl, #[strum(serialize = "bn254")] diff --git a/cli/src/util.rs b/cli/src/util.rs index 701c033d8..9a33f0bd0 100644 --- a/cli/src/util.rs +++ b/cli/src/util.rs @@ -13,6 +13,7 @@ macro_rules! clap_enum_variants { macro_rules! call_with_field { ($function:ident::<$field:ident>($($args:expr),*) ) => { match $field { + FieldArgument::Bb => $function::($($args),*), FieldArgument::Gl => $function::($($args),*), FieldArgument::Bn254 => $function::($($args),*), } diff --git a/number/Cargo.toml b/number/Cargo.toml index 989adac51..cc73cdd29 100644 --- a/number/Cargo.toml +++ b/number/Cargo.toml @@ -9,17 +9,23 @@ repository = { workspace = true } [dependencies] ark-bn254 = { version = "0.4.0", default-features = false, features = [ - "scalar_field", + "scalar_field", ] } ark-ff = "0.4.2" ark-serialize = "0.4.2" +p3-baby-bear = "0.1.0" +p3-field = "0.1.0" num-bigint = { version = "0.4.3", features = ["serde"] } num-traits = "0.2.15" csv = "1.3" -serde = { version = "1.0", default-features = false, features = ["alloc", "derive", "rc"] } +serde = { version = "1.0", default-features = false, features = [ + "alloc", + "derive", + "rc", +] } serde_with = "3.6.1" -schemars = { version = "0.8.16", features = ["preserve_order"]} -ibig = { version = "0.3.6", features = ["serde"]} +schemars = { version = "0.8.16", features = ["preserve_order"] } +ibig = { version = "0.3.6", features = ["serde"] } serde_cbor = "0.11.2" derive_more = "0.99.17" diff --git a/number/src/baby_bear.rs b/number/src/baby_bear.rs new file mode 100644 index 000000000..94a75d017 --- /dev/null +++ b/number/src/baby_bear.rs @@ -0,0 +1,467 @@ +use p3_baby_bear::BabyBear; +use schemars::{ + schema::{Schema, SchemaObject}, + JsonSchema, +}; +use serde::{Deserialize, Serialize}; + +use num_traits::{ConstOne, ConstZero}; +use std::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Not, Sub, SubAssign}; +use std::str::FromStr; +use std::{collections::BTreeSet, fmt::LowerHex}; + +use crate::{BigUint, FieldElement, KnownField, LargeInt}; +use ark_ff::{One, Zero}; + +use core::fmt::{self, Debug, Formatter}; +use core::hash::Hash; + +use p3_field::{AbstractField, Field, PrimeField32}; + +#[derive( + Debug, + Copy, + Clone, + Default, + Eq, + Hash, + PartialEq, + Ord, + PartialOrd, + Serialize, + Deserialize, + derive_more::Display, +)] +pub struct BabyBearField(BabyBear); + +const P: u32 = 0x78000001; +const MONTY_BITS: u32 = 32; +const MONTY_ZERO: u32 = to_monty(0); +const MONTY_ONE: u32 = to_monty(1); + +const fn to_monty(x: u32) -> u32 { + (((x as u64) << MONTY_BITS) % P as u64) as u32 +} + +impl BabyBearField { + const ORDER: u32 = P; + + #[inline(always)] + fn from_canonical_u32(n: u32) -> Self { + Self(BabyBear::from_canonical_u32(n)) + } + + #[inline] + fn to_canonical_u32(self) -> u32 { + self.0.as_canonical_u32() + } +} + +impl FieldElement for BabyBearField { + type Integer = BBLargeInt; + + const BITS: u32 = 31; + + fn to_degree(&self) -> crate::DegreeType { + self.to_canonical_u32() as u64 + } + + fn to_integer(&self) -> Self::Integer { + self.to_canonical_u32().into() + } + + #[inline] + fn modulus() -> Self::Integer { + Self::ORDER.into() + } + + fn pow(self, exp: Self::Integer) -> Self { + Self(BabyBear::exp_u64_generic( + self.0, + exp.try_into_u64().unwrap(), + )) + } + + fn to_bytes_le(&self) -> Vec { + self.to_canonical_u32().to_le_bytes().to_vec() + } + + fn from_bytes_le(bytes: &[u8]) -> Self { + let u = u32::from_le_bytes(bytes.try_into().unwrap()); + Self::from_canonical_u32(u) + } + + fn from_str_radix(s: &str, radix: u32) -> Result { + u32::from_str_radix(s, radix) + .map(Self::from_canonical_u32) + .map_err(|e| e.to_string()) + } + + fn checked_from(value: ibig::UBig) -> Option { + if value < Self::modulus().to_arbitrary_integer() { + Some(u32::try_from(value).unwrap().into()) + } else { + None + } + } + + fn is_in_lower_half(&self) -> bool { + self.to_canonical_u32() <= (Self::ORDER - 1) / 2 + } + + fn known_field() -> Option { + Some(KnownField::BabyBearField) + } + + fn try_into_i32(&self) -> Option { + Some(self.to_canonical_u32() as i32) + } +} + +impl LowerHex for BabyBearField { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + LowerHex::fmt(&self.to_canonical_u32(), f) + } +} + +impl From for BabyBearField { + fn from(b: bool) -> Self { + Self(BabyBear::from_bool(b)) + } +} + +impl From for BabyBearField { + fn from(n: i64) -> Self { + From::::from(n as u64) + } +} + +impl From for BabyBearField { + fn from(n: i32) -> Self { + From::::from(n as i64) + } +} + +impl From for BabyBearField { + fn from(n: u32) -> Self { + Self(BabyBear::from_wrapped_u32(n)) + } +} + +impl From for BabyBearField { + #[inline] + fn from(n: u64) -> Self { + Self(BabyBear::from_wrapped_u64(n)) + } +} + +impl From for BabyBearField { + fn from(n: crate::BigUint) -> Self { + u64::try_from(n).unwrap().into() + } +} + +impl From for BabyBearField { + #[inline] + fn from(n: BBLargeInt) -> Self { + n.0.into() + } +} + +impl ConstZero for BabyBearField { + // TODO This is a hack because BabyBear::new() is private. + const ZERO: Self = unsafe { Self(std::mem::transmute(MONTY_ZERO)) }; +} + +impl Zero for BabyBearField { + fn zero() -> Self { + Self(BabyBear::zero()) + } + + fn is_zero(&self) -> bool { + self.0.is_zero() + } +} + +impl ConstOne for BabyBearField { + // TODO This is a hack because BabyBear::new() is private. + const ONE: Self = unsafe { Self(std::mem::transmute(MONTY_ONE)) }; +} + +impl One for BabyBearField { + fn one() -> Self { + Self(BabyBear::one()) + } + + fn is_one(&self) -> bool { + self.to_canonical_u32() == 1 + } +} + +impl FromStr for BabyBearField { + type Err = String; + fn from_str(s: &str) -> Result { + let n = BigUint::from_str(s).map_err(|e| e.to_string())?; + let modulus = Self::modulus(); + if n >= modulus.to_arbitrary_integer() { + Err(format!("Decimal number \"{s}\" too large for field.")) + } else { + Ok(n.into()) + } + } +} + +impl Neg for BabyBearField { + type Output = Self; + + #[inline] + fn neg(self) -> Self { + Self(self.0.neg()) + } +} + +impl Add for BabyBearField { + type Output = Self; + + #[inline] + fn add(self, rhs: Self) -> Self { + Self(self.0.add(rhs.0)) + } +} + +impl AddAssign for BabyBearField { + #[inline] + fn add_assign(&mut self, rhs: Self) { + self.0.add_assign(rhs.0) + } +} + +impl Sub for BabyBearField { + type Output = Self; + + #[inline] + fn sub(self, rhs: Self) -> Self { + Self(self.0.sub(rhs.0)) + } +} + +impl SubAssign for BabyBearField { + #[inline] + fn sub_assign(&mut self, rhs: Self) { + self.0.sub_assign(rhs.0) + } +} + +impl Mul for BabyBearField { + type Output = Self; + + fn mul(self, rhs: Self) -> Self { + Self(self.0.mul(rhs.0)) + } +} + +impl MulAssign for BabyBearField { + fn mul_assign(&mut self, rhs: Self) { + self.0.mul_assign(rhs.0) + } +} + +impl Div for BabyBearField { + type Output = Self; + + fn div(self, rhs: Self) -> Self::Output { + Self(self.0.div(rhs.0)) + } +} + +impl JsonSchema for BabyBearField { + fn schema_name() -> String { + "BabyBearField".to_string() + } + + fn json_schema(gen: &mut schemars::gen::SchemaGenerator) -> Schema { + // Since BabyBear is just a wrapper around u32, use the schema for u32 + let u32_schema = gen.subschema_for::(); + + SchemaObject { + // Define the schema for BabyBearField, where field is of type BabyBear (which is u32) + instance_type: Some(schemars::schema::InstanceType::Object.into()), + object: Some(Box::new(schemars::schema::ObjectValidation { + properties: vec![("field".to_string(), u32_schema)] + .into_iter() + .collect(), + required: BTreeSet::from(["field".to_string()]), // Convert Vec to BTreeSet + ..Default::default() + })), + ..Default::default() + } + .into() + } +} + +#[derive( + Clone, + Copy, + PartialEq, + Eq, + Debug, + Default, + PartialOrd, + Ord, + Hash, + derive_more::Display, + Serialize, + Deserialize, + JsonSchema, + derive_more::Mul, + derive_more::Add, + derive_more::Sub, + derive_more::AddAssign, + derive_more::SubAssign, + derive_more::MulAssign, + derive_more::Shr, + derive_more::Shl, + derive_more::BitAnd, + derive_more::BitOr, + derive_more::BitXor, + derive_more::BitAndAssign, + derive_more::BitOrAssign, + derive_more::BitXorAssign, +)] +pub struct BBLargeInt(u32); + +impl LargeInt for BBLargeInt { + const NUM_BITS: usize = 32; + + fn to_arbitrary_integer(self) -> ibig::UBig { + self.0.into() + } + + fn num_bits(&self) -> usize { + Self::NUM_BITS - self.0.leading_zeros() as usize + } + + fn one() -> Self { + Self(1) + } + + fn is_one(&self) -> bool { + self.0 == 1 + } + + fn try_into_u64(&self) -> Option { + Some(self.0 as u64) + } + + fn try_into_u32(&self) -> Option { + Some(self.0) + } + + fn from_hex(s: &str) -> Self { + Self(u32::from_str_radix(s, 16).unwrap()) + } +} + +impl From for BBLargeInt { + fn from(value: u32) -> Self { + Self(value) + } +} + +impl From for BBLargeInt { + fn from(value: u64) -> Self { + Self(value as u32) + } +} + +impl Zero for BBLargeInt { + fn zero() -> Self { + Self(0) + } + + fn is_zero(&self) -> bool { + self.0 == 0 + } +} + +impl ConstZero for BBLargeInt { + const ZERO: Self = Self(0); +} + +impl LowerHex for BBLargeInt { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + LowerHex::fmt(&self.0, f) + } +} + +impl Not for BBLargeInt { + type Output = Self; + + fn not(self) -> Self::Output { + Self(!self.0) + } +} + +#[cfg(test)] +mod test { + use crate::traits::int_from_hex_str; + use test_log::test; + + use super::*; + + #[test] + fn bitwise() { + let n = int_from_hex_str::("00ff00ff"); + let p = int_from_hex_str::("f00ff00f"); + let not_n = int_from_hex_str::("ff00ff00"); + let n_shr_4 = int_from_hex_str::("000ff00f"); + let n_shl_4 = int_from_hex_str::("0ff00ff0"); + let n_or_p = int_from_hex_str::("f0fff0ff"); + let n_and_p = int_from_hex_str::("000f000f"); + let n_xor_p = int_from_hex_str::("f0f0f0f0"); + + assert_eq!(n.not().not(), n); + assert_eq!(n.not(), not_n); + assert_eq!(n >> 4, n_shr_4); + assert_eq!(n << 4, n_shl_4); + assert_eq!(n & p, n_and_p); + assert_eq!(n | p, n_or_p); + assert_eq!(n ^ p, n_xor_p); + } + + #[test] + fn zero_one() { + let x = BabyBearField::ZERO; + assert_eq!(x, BabyBearField::zero()); + assert_eq!(x.to_canonical_u32(), 0); + let y = BabyBearField::ONE; + assert_eq!(y, BabyBearField::one()); + assert_eq!(y.to_canonical_u32(), 1); + let z = x + y + y; + assert_eq!(z.to_canonical_u32(), 2); + } + + #[test] + fn lower_half() { + let x = BabyBearField::from(0); + assert!(x.is_in_lower_half()); + assert!(!(x - 1.into()).is_in_lower_half()); + + let y = BabyBearField::from_str_radix("3c000000", 16).unwrap(); + assert!(y.is_in_lower_half()); + assert!(!(y + 1.into()).is_in_lower_half()); + } + + #[test] + #[should_panic] + fn integer_div_by_zero() { + let _ = BabyBearField::from(1).to_arbitrary_integer() + / BabyBearField::from(0).to_arbitrary_integer(); + } + + #[test] + #[should_panic] + fn div_by_zero() { + let _ = BabyBearField::from(1) / BabyBearField::from(0); + } +} diff --git a/number/src/lib.rs b/number/src/lib.rs index 8842281c6..5634ed9c7 100644 --- a/number/src/lib.rs +++ b/number/src/lib.rs @@ -4,6 +4,7 @@ #[macro_use] mod macros; +mod baby_bear; mod bn254; mod goldilocks; mod serialize; @@ -12,6 +13,7 @@ pub use serialize::{ buffered_write_file, read_polys_csv_file, write_polys_csv_file, CsvRenderMode, ReadWrite, }; +pub use baby_bear::BabyBearField; pub use bn254::Bn254Field; pub use goldilocks::GoldilocksField; pub use traits::KnownField; diff --git a/number/src/traits.rs b/number/src/traits.rs index 34eca3071..9e894ed6a 100644 --- a/number/src/traits.rs +++ b/number/src/traits.rs @@ -64,6 +64,7 @@ pub trait LargeInt: #[derive(Debug, PartialEq, Serialize, Deserialize)] pub enum KnownField { + BabyBearField, GoldilocksField, Bn254Field, } diff --git a/riscv/src/code_gen.rs b/riscv/src/code_gen.rs index ce26ad3df..c5929a69d 100644 --- a/riscv/src/code_gen.rs +++ b/riscv/src/code_gen.rs @@ -683,6 +683,7 @@ fn mul_instruction(runtime: &Runtime) -> &'static str { link ~> regs.mstore(W, STEP + 3, tmp4_col); "# } + KnownField::BabyBearField => todo!(), } }