soft3/hemera/rs/src/field.rs

// ---
// tags: hemera, rust
// crystal-type: source
// crystal-domain: comp
// ---
//! Goldilocks prime field (p = 2^64 - 2^32 + 1).
//!
//! Minimal implementation covering the operations hemera needs:
//! addition, subtraction, multiplication, and the x^7 S-box.

use core::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};

/// The Goldilocks prime: p = 2^64 - 2^32 + 1.
pub const P: u64 = 0xFFFF_FFFF_0000_0001;

/// Two's complement of P modulo 2^64: 2^32 - 1.
const NEG_ORDER: u64 = P.wrapping_neg(); // 0xFFFF_FFFF

/// A Goldilocks field element.
///
/// Internal value may be non-canonical (in `[0, 2^64)`).
/// Use `as_canonical_u64()` to reduce to `[0, p)`.
#[derive(Clone, Copy, Default, PartialEq, Eq, Hash)]
#[repr(transparent)]
pub struct Goldilocks {
    value: u64,
}

impl Goldilocks {
    pub const ZERO: Self = Self { value: 0 };

    #[inline]
    pub const fn new(value: u64) -> Self {
        Self { value }
    }

    /// Reduce to canonical form in [0, p).
    #[inline]
    pub fn as_canonical_u64(self) -> u64 {
        let mut c = self.value;
        if c >= P {
            c -= P;
        }
        c
    }

    /// Compute x^2.
    #[inline]
    fn square(self) -> Self {
        self * self
    }

    /// Compute x^7 (the Poseidon2 full-round S-box for Goldilocks).
    #[inline]
    pub fn pow7(self) -> Self {
        let x2 = self.square();
        let x3 = x2 * self;
        let x4 = x2.square();
        x3 * x4
    }

    /// Compute x^(-1) = x^(p-2) (the Poseidon2 partial-round S-box).
    ///
    /// Uses Fermat's little theorem: x^(-1) = x^(p-2) mod p.
    /// Convention: 0^(-1) = 0 (permutation over F_p).
    #[inline]
    pub fn inv(self) -> Self {
        // x^(p-2) by the 75-multiply addition chain (acpu gl_inv) instead
        // of naive square-and-multiply (~125 serial muls). The partial
        // rounds of the Poseidon2 permutation call this once each, and the
        // permutation is the single hottest primitive in the stack — every
        // saved serial multiply here is felt from Brakedown to the graph.
        // Bit-identical result: same Fermat exponent, better chain.
        if self.as_canonical_u64() == 0 {
            return Self::ZERO;
        }
        let x = self;
        // chain of 2^k-1 powers up to x^(2^31-1)
        let x2 = x * x;
        let x3 = x2 * x;
        let x4 = x2 * x2;
        let x7 = x3 * x4;
        let x6 = x3 * x3;
        let x12 = x6 * x6;
        let x15 = x12 * x3; // 2^4-1
        let x30 = x15 * x15;
        let x60 = x30 * x30;
        let x120 = x60 * x60;
        let x127 = x120 * x7; // 2^7-1
        let x254 = x127 * x127;
        let x255 = x254 * x; // 2^8-1
        let mut t = x255;
        for _ in 0..7 {
            t = t * t;
        }
        let x_2p15m1 = t * x127; // 2^15-1
        let x_2p16m2 = x_2p15m1 * x_2p15m1;
        let x_2p16m1 = x_2p16m2 * x; // 2^16-1
        t = x_2p16m1;
        for _ in 0..15 {
            t = t * t;
        }
        let x_2p31m1 = t * x_2p15m1; // 2^31-1
        // x^(2^32-1)
        let x_2p32m2 = x_2p31m1 * x_2p31m1;
        let x_epsilon = x_2p32m2 * x;
        // x^((2^31-1)*2^33)
        t = x_2p31m1;
        for _ in 0..33 {
            t = t * t;
        }
        // p-2 = (2^31-1)*2^33 + 2^32-1
        t * x_epsilon
    }

    /// Double this element.
    #[inline]
    fn double(self) -> Self {
        self + self
    }
}

impl core::fmt::Debug for Goldilocks {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "Goldilocks({})", self.as_canonical_u64())
    }
}

// ── Arithmetic ──────────────────────────────────────────────────────

impl Add for Goldilocks {
    type Output = Self;

    #[inline]
    fn add(self, rhs: Self) -> Self {
        let (sum, over) = self.value.overflowing_add(rhs.value);
        let (mut sum, over) = sum.overflowing_add(u64::from(over) * NEG_ORDER);
        if over {
            sum += NEG_ORDER;
        }
        Self::new(sum)
    }
}

impl AddAssign for Goldilocks {
    #[inline]
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl Sub for Goldilocks {
    type Output = Self;

    #[inline]
    fn sub(self, rhs: Self) -> Self {
        let (diff, under) = self.value.overflowing_sub(rhs.value);
        let (mut diff, under) = diff.overflowing_sub(u64::from(under) * NEG_ORDER);
        if under {
            diff -= NEG_ORDER;
        }
        Self::new(diff)
    }
}

impl SubAssign for Goldilocks {
    #[inline]
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl Mul for Goldilocks {
    type Output = Self;

    #[inline]
    fn mul(self, rhs: Self) -> Self {
        reduce128(u128::from(self.value) * u128::from(rhs.value))
    }
}

impl MulAssign for Goldilocks {
    #[inline]
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl core::iter::Sum for Goldilocks {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        reduce128(iter.map(|x| x.value as u128).sum::<u128>())
    }
}

/// Reduce a 128-bit product to a Goldilocks element.
///
/// Uses the identity 2^64 ≡ 2^32 - 1 (mod p).
#[inline]
fn reduce128(x: u128) -> Goldilocks {
    let x_lo = x as u64;
    let x_hi = (x >> 64) as u64;
    let x_hi_hi = x_hi >> 32;
    let x_hi_lo = x_hi & NEG_ORDER;

    let (mut t0, borrow) = x_lo.overflowing_sub(x_hi_hi);
    if borrow {
        t0 -= NEG_ORDER;
    }
    let t1 = x_hi_lo * NEG_ORDER;
    let (res, carry) = t0.overflowing_add(t1);
    Goldilocks::new(res + NEG_ORDER * u64::from(carry))
}

// ── MDS and diffusion matrices ──────────────────────────────────────

/// Apply the 4x4 MDS matrix used in Poseidon2 external rounds:
/// ```text
/// [ 2 3 1 1 ]
/// [ 1 2 3 1 ]
/// [ 1 1 2 3 ]
/// [ 3 1 1 2 ]
/// ```
#[inline(always)]
pub fn apply_mat4(x: &mut [Goldilocks; 4]) {
    let t01 = x[0] + x[1];
    let t23 = x[2] + x[3];
    let t0123 = t01 + t23;
    let t01123 = t0123 + x[1];
    let t01233 = t0123 + x[3];
    x[3] = t01233 + x[0].double();
    x[1] = t01123 + x[2].double();
    x[0] = t01123 + t01;
    x[2] = t01233 + t23;
}

/// Apply the external MDS layer for a width-16 state.
///
/// Multiplies by the 16×16 circulant-of-4×4 matrix:
/// `[[2M M ... M], [M 2M ... M], ..., [M M ... 2M]]`
/// where M is the 4×4 MDS matrix.
#[inline]
pub fn mds_light_permutation(state: &mut [Goldilocks; 16]) {
    // Apply M4 to each consecutive 4-element chunk.
    for chunk in state.chunks_exact_mut(4) {
        apply_mat4(chunk.try_into().unwrap());
    }

    // Compute column sums (one per M4 column position).
    let sums: [Goldilocks; 4] = core::array::from_fn(|k| {
        (0..16).step_by(4).map(|j| state[j + k]).sum()
    });

    // Add the appropriate column sum to each element.
    for (i, elem) in state.iter_mut().enumerate() {
        *elem += sums[i % 4];
    }
}

/// Diagonal elements of the internal diffusion matrix for Goldilocks t=16.
///
/// The full matrix is M_I = 1 + diag(d), where 1 is the all-ones matrix.
pub const MATRIX_DIAG_16: [Goldilocks; 16] = [
    Goldilocks::new(0xde9b91a467d6afc0),
    Goldilocks::new(0xc5f16b9c76a9be17),
    Goldilocks::new(0x0ab0fef2d540ac55),
    Goldilocks::new(0x3001d27009d05773),
    Goldilocks::new(0xed23b1f906d3d9eb),
    Goldilocks::new(0x5ce73743cba97054),
    Goldilocks::new(0x1c3bab944af4ba24),
    Goldilocks::new(0x2faa105854dbafae),
    Goldilocks::new(0x53ffb3ae6d421a10),
    Goldilocks::new(0xbcda9df8884ba396),
    Goldilocks::new(0xfc1273e4a31807bb),
    Goldilocks::new(0xc77952573d5142c0),
    Goldilocks::new(0x56683339a819b85e),
    Goldilocks::new(0x328fcbd8f0ddc8eb),
    Goldilocks::new(0xb5101e303fce9cb7),
    Goldilocks::new(0x774487b8c40089bb),
];

/// Apply the internal diffusion matrix: M_I = 1 + diag(d).
///
/// Computes `state'[i] = d[i] * state[i] + sum(state)`.
#[inline]
pub fn matmul_internal(state: &mut [Goldilocks; 16]) {
    let sum: Goldilocks = state.iter().copied().sum();
    for i in 0..16 {
        state[i] *= MATRIX_DIAG_16[i];
        state[i] += sum;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The addition chain must equal naive x^(p-2) everywhere — checked
    /// against x*inv(x)==1 over a spread of values including edge cases.
    #[test]
    fn inv_chain_matches_fermat() {
        let mut v = [0u64; 72];
        v[..8].copy_from_slice(&[1, 2, 3, P - 1, P - 2, 0xFFFF_FFFF, 1 << 32, u64::MAX % P]);
        let mut x = 0x9e3779b97f4a7c15u64;
        for slot in v[8..].iter_mut() {
            x = x.wrapping_mul(0xbf58476d1ce4e5b9).rotate_left(31);
            *slot = x % P;
        }
        for &raw in &v {
            if raw == 0 {
                continue;
            }
            let a = Goldilocks::new(raw);
            assert_eq!(
                (a * a.inv()).as_canonical_u64(),
                1,
                "inv broken at {raw}"
            );
        }
        assert_eq!(Goldilocks::ZERO.inv().as_canonical_u64(), 0);
    }

    use super::*;

    #[test]
    fn field_add_basic() {
        let a = Goldilocks::new(1);
        let b = Goldilocks::new(2);
        assert_eq!((a + b).as_canonical_u64(), 3);
    }

    #[test]
    fn field_add_wrap() {
        let a = Goldilocks::new(P - 1);
        let b = Goldilocks::new(1);
        assert_eq!((a + b).as_canonical_u64(), 0);
    }

    #[test]
    fn field_sub_basic() {
        let a = Goldilocks::new(5);
        let b = Goldilocks::new(3);
        assert_eq!((a - b).as_canonical_u64(), 2);
    }

    #[test]
    fn field_sub_wrap() {
        let a = Goldilocks::new(0);
        let b = Goldilocks::new(1);
        assert_eq!((a - b).as_canonical_u64(), P - 1);
    }

    #[test]
    fn field_mul_basic() {
        let a = Goldilocks::new(3);
        let b = Goldilocks::new(7);
        assert_eq!((a * b).as_canonical_u64(), 21);
    }

    #[test]
    fn field_mul_large() {
        // (P-1) * (P-1) mod P = 1
        let a = Goldilocks::new(P - 1);
        assert_eq!((a * a).as_canonical_u64(), 1);
    }

    #[test]
    fn pow7_basic() {
        let x = Goldilocks::new(2);
        assert_eq!(x.pow7().as_canonical_u64(), 128);
    }

    #[test]
    fn pow7_zero() {
        assert_eq!(Goldilocks::ZERO.pow7().as_canonical_u64(), 0);
    }

    #[test]
    fn pow7_one() {
        assert_eq!(Goldilocks::new(1).pow7().as_canonical_u64(), 1);
    }

    #[test]
    fn inv_basic() {
        // 2 * inv(2) = 1
        let x = Goldilocks::new(2);
        let x_inv = x.inv();
        assert_eq!((x * x_inv).as_canonical_u64(), 1);
    }

    #[test]
    fn inv_one() {
        assert_eq!(Goldilocks::new(1).inv().as_canonical_u64(), 1);
    }

    #[test]
    fn inv_zero() {
        assert_eq!(Goldilocks::ZERO.inv().as_canonical_u64(), 0);
    }

    #[test]
    fn inv_roundtrip() {
        let x = Goldilocks::new(42);
        let x_inv = x.inv();
        assert_eq!((x * x_inv).as_canonical_u64(), 1);
        assert_eq!(x_inv.inv().as_canonical_u64(), x.as_canonical_u64());
    }

    #[test]
    fn inv_p_minus_one() {
        // (p-1)^(-1) = p-1 since (p-1)^2 = 1 mod p
        let a = Goldilocks::new(P - 1);
        assert_eq!(a.inv().as_canonical_u64(), P - 1);
    }

    #[test]
    fn canonical_reduces() {
        // A non-canonical value >= P
        let a = Goldilocks::new(P);
        assert_eq!(a.as_canonical_u64(), 0);
        let b = Goldilocks::new(P + 1);
        assert_eq!(b.as_canonical_u64(), 1);
    }
}