Extract stable stwo-types crate containing field and struct types used downstream (#1331)

Author: Arcticae

Cargo.lock

@@ -6,6 +6,7 @@ members = [
     "crates/constraint-framework",
     "crates/examples",
     "crates/std-shims",
+    "crates/stwo-types",
 ]
 exclude = ["ensure-verifier-no_std"]
 resolver = "2"
@@ -32,6 +33,7 @@ rayon = { version = "1.10.0", optional = false }
 rand = { version = "0.8.5", default-features = false, features = ["small_rng"] }
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 hashbrown = { version = ">=0.15.2", features = ["serde"] }
+stwo-types = { path = "crates/stwo-types", version = "2.0.1", default-features = false }
 
 [profile.bench]
 codegen-units = 1

Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "stwo-types"
+version = "2.0.1"
+edition.workspace = true
+license.workspace = true
+description = "Core field types shared across stwo crates"
+
+[features]
+default = ["std"]
+std = ["serde/std", "num-traits/std", "rand/std", "std-shims/std"]
+parallel = ["dep:rayon"]
+
+[dependencies]
+bytemuck = { workspace = true, features = ["derive", "extern_crate_alloc"] }
+num-traits.workspace = true
+rand.workspace = true
+serde.workspace = true
+rayon = { workspace = true, optional = true }
+std-shims = { path = "../std-shims", version = "1.0.0", default-features = false }
+
+[lints.rust]
+future-incompatible = "deny"
+nonstandard-style = "deny"
+rust-2018-idioms = "deny"
+
+[lints.clippy]
+missing_const_for_fn = "warn"

crates/stwo-types/Cargo.toml

@@ -0,0 +1,16 @@
+use serde::{Deserialize, Serialize};
+
+use super::m31::M31;
+
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, Default, Eq, PartialEq, Hash)]
+pub struct CasmState {
+    pub pc: M31,
+    pub ap: M31,
+    pub fp: M31,
+}
+
+impl CasmState {
+    pub const fn values(&self) -> [M31; 3] {
+        [self.pc, self.ap, self.fp]
+    }
+}

crates/stwo-types/src/fields/casm_state.rs

@@ -0,0 +1,115 @@
+use core::fmt::{Debug, Display};
+use core::ops::{
+    Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Sub, SubAssign,
+};
+
+use serde::{Deserialize, Serialize};
+
+use super::m31::M31;
+use super::{ComplexConjugate, FieldExpOps};
+use crate::{impl_extension_field, impl_field};
+pub const P2: u64 = 4611686014132420609; // (2 ** 31 - 1) ** 2
+
+/// Complex extension field of M31.
+/// Equivalent to M31\[x\] over (x^2 + 1) as the irreducible polynomial.
+/// Represented as (a, b) of a + bi.
+#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
+pub struct CM31(pub M31, pub M31);
+
+impl_field!(CM31, P2);
+impl_extension_field!(CM31, M31);
+
+impl CM31 {
+    pub const fn from_u32_unchecked(a: u32, b: u32) -> CM31 {
+        Self(M31::from_u32_unchecked(a), M31::from_u32_unchecked(b))
+    }
+
+    pub const fn from_m31(a: M31, b: M31) -> CM31 {
+        Self(a, b)
+    }
+}
+
+impl Display for CM31 {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "{} + {}i", self.0, self.1)
+    }
+}
+
+impl Debug for CM31 {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        write!(f, "{} + {}i", self.0, self.1)
+    }
+}
+
+impl Mul for CM31 {
+    type Output = Self;
+
+    fn mul(self, rhs: Self) -> Self::Output {
+        // (a + bi) * (c + di) = (ac - bd) + (ad + bc)i.
+        Self(
+            self.0 * rhs.0 - self.1 * rhs.1,
+            self.0 * rhs.1 + self.1 * rhs.0,
+        )
+    }
+}
+
+impl TryInto<M31> for CM31 {
+    type Error = ();
+
+    fn try_into(self) -> Result<M31, Self::Error> {
+        if self.1 != M31::zero() {
+            return Err(());
+        }
+        Ok(self.0)
+    }
+}
+
+impl FieldExpOps for CM31 {
+    fn inverse(&self) -> Self {
+        assert!(!self.is_zero(), "0 has no inverse");
+        // 1 / (a + bi) = (a - bi) / (a^2 + b^2).
+        Self(self.0, -self.1) * (self.0.square() + self.1.square()).inverse()
+    }
+}
+
+#[cfg(test)]
+#[macro_export]
+macro_rules! cm31 {
+    ($m0:expr, $m1:expr) => {
+        CM31::from_u32_unchecked($m0, $m1)
+    };
+}
+
+#[cfg(test)]
+mod tests {
+    use super::CM31;
+    use crate::fields::m31::P;
+    use crate::fields::FieldExpOps;
+    use crate::m31;
+
+    #[test]
+    fn test_inverse() {
+        let cm = cm31!(1, 2);
+        let cm_inv = cm.inverse();
+        assert_eq!(cm * cm_inv, cm31!(1, 0));
+    }
+
+    #[test]
+    fn test_ops() {
+        let cm0 = cm31!(1, 2);
+        let cm1 = cm31!(4, 5);
+        let m = m31!(8);
+        let cm = CM31::from(m);
+        let cm0_x_cm1 = cm31!(P - 6, 13);
+
+        assert_eq!(cm0 + cm1, cm31!(5, 7));
+        assert_eq!(cm1 + m, cm1 + cm);
+        assert_eq!(cm0 * cm1, cm0_x_cm1);
+        assert_eq!(cm1 * m, cm1 * cm);
+        assert_eq!(-cm0, cm31!(P - 1, P - 2));
+        assert_eq!(cm0 - cm1, cm31!(P - 3, P - 3));
+        assert_eq!(cm1 - m, cm1 - cm);
+        assert_eq!(cm0_x_cm1 / cm1, cm31!(1, 2));
+        assert_eq!(cm1 / m, cm1 / cm);
+    }
+}