tokitai-operator 0.1.0

//! Finite field domain and element type.
//!
//! Supports prime fields F_p and prime-power extension fields
//! F_{p^k}. Extension fields are represented as polynomials in
//! F_p[x] modulo a chosen irreducible polynomial of degree k; the
//! polynomial is provided by the caller (P346 acceptance tests use
//! the canonical x^2 + x + 1 over F_2 for the k=2 case).
//!
//! P346 is a domain-completeness pass: the element type, the Domain
//! trait implementation, and the FFE claim (prime, extension_degree)
//! are added. No GPU lowering, no Galois group computation, no
//! polynomial ring.

use std::fmt;
use std::ops::{Add, Mul, Neg, Sub};

use serde::{Deserialize, Serialize};

use super::contract::{Claim, ContractSet, Scope};
use super::precision::PrecisionModel;
use super::structure::AlgebraicStructure;
use super::{Domain, DomainId};
use crate::{Error, Result};

/// Element of a finite field `F_{p^k}`.
///
/// For prime fields (k=1), the underlying value is just the residue
/// 0..p. For extension fields (k>1), values are polynomials in
/// `F_p[x]` modulo an irreducible polynomial of degree k; the
/// polynomial is held as a `Vec<u64>` of coefficients (length k,
/// little-endian) and the modulus polynomial is held by the
/// `FiniteFieldDomain`.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FiniteFieldElement {
    /// Element of a prime field F_p (extension_degree = 1).
    Prime { residue: u64, prime: u64 },
    /// Element of an extension field F_{p^k}, k >= 2.
    /// Coefficients are little-endian (constant term first).
    Extension { coeffs: Vec<u64>, prime: u64 },
}

impl FiniteFieldElement {
    /// Construct the additive identity.
    pub fn zero(prime: u64, extension_degree: usize) -> Self {
        if extension_degree <= 1 {
            FiniteFieldElement::Prime { residue: 0, prime }
        } else {
            FiniteFieldElement::Extension {
                coeffs: vec![0; extension_degree],
                prime,
            }
        }
    }

    /// Construct the multiplicative identity.
    pub fn one(prime: u64, extension_degree: usize) -> Self {
        if extension_degree <= 1 {
            FiniteFieldElement::Prime {
                residue: 1 % prime,
                prime,
            }
        } else {
            let mut coeffs = vec![0; extension_degree];
            coeffs[0] = 1 % prime;
            FiniteFieldElement::Extension { coeffs, prime }
        }
    }

    /// Return the underlying prime.
    pub fn prime(&self) -> u64 {
        match self {
            FiniteFieldElement::Prime { prime, .. } => *prime,
            FiniteFieldElement::Extension { prime, .. } => *prime,
        }
    }

    /// Return the extension degree (1 for prime fields).
    pub fn extension_degree(&self) -> usize {
        match self {
            FiniteFieldElement::Prime { .. } => 1,
            FiniteFieldElement::Extension { coeffs, .. } => coeffs.len(),
        }
    }

    /// Add two elements (modular reduction).
    pub fn add(&self, rhs: &FiniteFieldElement) -> Result<FiniteFieldElement> {
        match (self, rhs) {
            (
                FiniteFieldElement::Prime {
                    residue: a,
                    prime: p,
                },
                FiniteFieldElement::Prime {
                    residue: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                Ok(FiniteFieldElement::Prime {
                    residue: (a + b) % p,
                    prime: *p,
                })
            }
            (
                FiniteFieldElement::Extension {
                    coeffs: a,
                    prime: p,
                },
                FiniteFieldElement::Extension {
                    coeffs: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                debug_assert_eq!(a.len(), b.len());
                let coeffs = a.iter().zip(b.iter()).map(|(x, y)| (x + y) % p).collect();
                Ok(FiniteFieldElement::Extension { coeffs, prime: *p })
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldElement::add: kind mismatch: self={self:?}, rhs={rhs:?} \
                 (add a Prime and an Extension element directly is not supported; \
                 use FiniteFieldDomain::add to dispatch over a known domain)"
            ))),
        }
    }

    /// Subtract two elements.
    pub fn sub(&self, rhs: &FiniteFieldElement) -> Result<FiniteFieldElement> {
        match (self, rhs) {
            (
                FiniteFieldElement::Prime {
                    residue: a,
                    prime: p,
                },
                FiniteFieldElement::Prime {
                    residue: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                Ok(FiniteFieldElement::Prime {
                    residue: (a + p - b % p) % p,
                    prime: *p,
                })
            }
            (
                FiniteFieldElement::Extension {
                    coeffs: a,
                    prime: p,
                },
                FiniteFieldElement::Extension {
                    coeffs: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                debug_assert_eq!(a.len(), b.len());
                let coeffs = a
                    .iter()
                    .zip(b.iter())
                    .map(|(x, y)| (x + p - y % p) % p)
                    .collect();
                Ok(FiniteFieldElement::Extension { coeffs, prime: *p })
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldElement::sub: kind mismatch: self={self:?}, rhs={rhs:?} \
                 (subtract a Prime and an Extension element directly is not supported)"
            ))),
        }
    }

    /// Multiply two prime-field elements. (Extension-field multiplication
    /// requires the modulus polynomial and is handled in the domain.)
    pub fn mul_prime(&self, rhs: &FiniteFieldElement) -> Option<FiniteFieldElement> {
        match (self, rhs) {
            (
                FiniteFieldElement::Prime {
                    residue: a,
                    prime: p,
                },
                FiniteFieldElement::Prime {
                    residue: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                Some(FiniteFieldElement::Prime {
                    residue: (a * b) % p,
                    prime: *p,
                })
            }
            _ => None,
        }
    }

    /// Negate an element.
    pub fn neg(&self) -> FiniteFieldElement {
        match self {
            FiniteFieldElement::Prime {
                residue: a,
                prime: p,
            } => FiniteFieldElement::Prime {
                residue: (p - a % p) % p,
                prime: *p,
            },
            FiniteFieldElement::Extension { coeffs, prime } => {
                let coeffs = coeffs.iter().map(|c| (prime - c % prime) % prime).collect();
                FiniteFieldElement::Extension {
                    coeffs,
                    prime: *prime,
                }
            }
        }
    }

    /// Fermat's little theorem: a^(p-1) == 1 for a in F_p*, a != 0.
    pub fn pow_prime(&self, mut exp: u64) -> Result<FiniteFieldElement> {
        match self {
            FiniteFieldElement::Prime {
                residue: a,
                prime: p,
            } => {
                if *a == 0 {
                    return Ok(FiniteFieldElement::Prime {
                        residue: 0,
                        prime: *p,
                    });
                }
                let mut base = *a;
                let mut result: u64 = 1;
                while exp > 0 {
                    if exp & 1 == 1 {
                        result = (result * base) % p;
                    }
                    exp >>= 1;
                    base = (base * base) % p;
                }
                Ok(FiniteFieldElement::Prime {
                    residue: result,
                    prime: *p,
                })
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldElement::pow_prime: called on extension element: self={self:?} \
                 (pow_prime is prime-field only; use the extension-domain dispatch for F_{{p^k}})"
            ))),
        }
    }

    /// Modular inverse in a prime field via Fermat's little theorem.
    pub fn inv_prime(&self) -> Result<FiniteFieldElement> {
        match self {
            FiniteFieldElement::Prime {
                residue: a,
                prime: p,
            } => {
                if *a == 0 {
                    return Err(Error::domain(format!(
                        "FiniteFieldElement::inv_prime: zero has no inverse (self={self:?})"
                    )));
                }
                self.pow_prime(p - 2)
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldElement::inv_prime: called on extension element: self={self:?} \
                 (inv_prime is prime-field only; for F_{{p^k}} inverses, route through the domain)"
            ))),
        }
    }

    /// Modular division in a prime field via Fermat's little theorem.
    pub fn div_prime(&self, rhs: &FiniteFieldElement) -> Result<FiniteFieldElement> {
        match (self, rhs) {
            (
                FiniteFieldElement::Prime {
                    residue: a,
                    prime: p,
                },
                FiniteFieldElement::Prime {
                    residue: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                if *b == 0 {
                    return Err(Error::domain(format!(
                        "FiniteFieldElement::div_prime: division by zero: \
                         self={self:?}, rhs={rhs:?}"
                    )));
                }
                // (a * b^(p-2)) mod p = a / b
                let mut result = 1u64;
                let mut base = *b;
                let mut exp = p - 2;
                while exp > 0 {
                    if exp & 1 == 1 {
                        result = (result * base) % p;
                    }
                    exp >>= 1;
                    base = (base * base) % p;
                }
                Ok(FiniteFieldElement::Prime {
                    residue: (a * result) % p,
                    prime: *p,
                })
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldElement::div_prime: called on non-prime element: self={self:?}, rhs={rhs:?} \
                 (div_prime is prime-field only; for F_{{p^k}} division, route through the domain)"
            ))),
        }
    }
}

// Extension-field polynomial multiplication mod irreducible poly.
// a, b have length k. modulus has length k+1 (the irreducible polynomial
// of degree k over F_p, with modulus[k] = 1 for monic).
// Result has length k.
fn poly_mul_mod(a: &[u64], b: &[u64], modulus: &[u64], prime: u64) -> Vec<u64> {
    let k = a.len();
    debug_assert_eq!(b.len(), k);
    debug_assert_eq!(modulus.len(), k + 1);
    let mut raw = vec![0u64; 2 * k];
    for i in 0..k {
        for j in 0..k {
            raw[i + j] = (raw[i + j] + a[i] * b[j]) % prime;
        }
    }
    // Reduce mod the modulus polynomial of degree k. For each
    // coefficient at position m >= k, replace x^m with
    // (-modulus[0..k]) * x^(m-k) and reduce.
    for m in (k..(2 * k)).rev() {
        if raw[m] == 0 {
            continue;
        }
        let coef = raw[m];
        raw[m] = 0;
        for i in 0..k {
            // raw[m - k + i] -= coef * modulus[i]  (mod prime)
            let v = (coef * modulus[i]) % prime;
            raw[m - k + i] = (raw[m - k + i] + prime - v) % prime;
        }
    }
    raw.truncate(k);
    raw
}

/// Domain for a finite field `F_{p^k}` with an explicit irreducible
/// polynomial of degree k (length k+1, little-endian coefficients).
/// For prime fields (k=1), the modulus is `[1, 0]` representing x.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FiniteFieldDomain {
    prime: u64,
    extension_degree: usize,
    /// Modulus polynomial coefficients, little-endian, length = extension_degree + 1.
    modulus: Vec<u64>,
    name: String,
}

impl FiniteFieldDomain {
    /// Construct a prime field F_p (extension_degree = 1).
    pub fn prime(prime: u64) -> Self {
        Self {
            prime,
            extension_degree: 1,
            modulus: vec![1, 0], // x
            name: format!("F_{}", prime),
        }
    }

    /// Construct an extension field F_{p^k} with the given modulus
    /// polynomial coefficients (little-endian, length = k+1).
    pub fn extension(prime: u64, extension_degree: usize, modulus: Vec<u64>) -> Self {
        assert_eq!(modulus.len(), extension_degree + 1);
        assert!(modulus[extension_degree] == 1, "modulus must be monic");
        Self {
            prime,
            extension_degree,
            modulus,
            name: format!("F_{}^{}", prime, extension_degree),
        }
    }

    pub fn prime_field(&self) -> u64 {
        self.prime
    }

    pub fn extension_degree(&self) -> usize {
        self.extension_degree
    }

    /// The zero element of this field.
    pub fn zero(&self) -> FiniteFieldElement {
        FiniteFieldElement::zero(self.prime, self.extension_degree)
    }

    /// The one element of this field.
    pub fn one(&self) -> FiniteFieldElement {
        FiniteFieldElement::one(self.prime, self.extension_degree)
    }

    /// Construct an element from raw little-endian coefficients.
    pub fn element(&self, coeffs: &[u64]) -> FiniteFieldElement {
        let n = self.extension_degree;
        let mut c = coeffs.to_vec();
        c.resize(n, 0);
        for v in &mut c {
            *v %= self.prime;
        }
        if n == 1 {
            FiniteFieldElement::Prime {
                residue: c[0],
                prime: self.prime,
            }
        } else {
            FiniteFieldElement::Extension {
                coeffs: c,
                prime: self.prime,
            }
        }
    }

    /// Add two elements in this domain.
    pub fn add(
        &self,
        lhs: &FiniteFieldElement,
        rhs: &FiniteFieldElement,
    ) -> Result<FiniteFieldElement> {
        lhs.add(rhs)
    }

    /// Subtract two elements.
    pub fn sub(
        &self,
        lhs: &FiniteFieldElement,
        rhs: &FiniteFieldElement,
    ) -> Result<FiniteFieldElement> {
        lhs.sub(rhs)
    }

    /// Multiply two elements. For extension fields, this is
    /// polynomial multiplication mod the modulus polynomial.
    pub fn mul(
        &self,
        lhs: &FiniteFieldElement,
        rhs: &FiniteFieldElement,
    ) -> Result<FiniteFieldElement> {
        match (lhs, rhs) {
            (
                FiniteFieldElement::Prime {
                    residue: a,
                    prime: p,
                },
                FiniteFieldElement::Prime {
                    residue: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                Ok(FiniteFieldElement::Prime {
                    residue: (a * b) % p,
                    prime: *p,
                })
            }
            (
                FiniteFieldElement::Extension {
                    coeffs: a,
                    prime: p,
                },
                FiniteFieldElement::Extension {
                    coeffs: b,
                    prime: q,
                },
            ) => {
                debug_assert_eq!(p, q);
                let reduced = poly_mul_mod(a, b, &self.modulus, *p);
                Ok(FiniteFieldElement::Extension {
                    coeffs: reduced,
                    prime: *p,
                })
            }
            _ => Err(Error::domain(format!(
                "FiniteFieldDomain::mul: kind mismatch: lhs={lhs:?}, rhs={rhs:?} \
                 (multiplied a Prime and an Extension element directly is not supported \
                 through the domain dispatch; build the Extension via the same prime as lhs)"
            ))),
        }
    }

    /// Negate an element.
    pub fn neg(&self, v: &FiniteFieldElement) -> FiniteFieldElement {
        v.neg()
    }
}

impl Domain for FiniteFieldDomain {
    type Element = FiniteFieldElement;

    fn id(&self) -> DomainId {
        DomainId::new(self.name.clone())
    }

    fn structure(&self) -> AlgebraicStructure {
        AlgebraicStructure::Field
    }

    fn precision_model(&self) -> PrecisionModel {
        PrecisionModel::Exact
    }

    fn contracts(&self) -> ContractSet {
        let scope = Scope::Domain(self.name.clone());
        ContractSet::from_iter([
            crate::domain::Contract::new(
                crate::domain::ContractId(80),
                Claim::Associative,
                scope.clone(),
                crate::domain::Evidence::Axiom,
            ),
            crate::domain::Contract::new(
                crate::domain::ContractId(81),
                Claim::Commutative,
                scope.clone(),
                crate::domain::Evidence::Axiom,
            ),
            crate::domain::Contract::new(
                crate::domain::ContractId(82),
                Claim::HasIdentity,
                scope.clone(),
                crate::domain::Evidence::Axiom,
            ),
            crate::domain::Contract::new(
                crate::domain::ContractId(83),
                Claim::HasInverse,
                scope.clone(),
                crate::domain::Evidence::Axiom,
            ),
            crate::domain::Contract::new(
                crate::domain::ContractId(84),
                Claim::Exact,
                scope.clone(),
                crate::domain::Evidence::Axiom,
            ),
            crate::domain::Contract::new(
                crate::domain::ContractId(85),
                Claim::FiniteFieldSpec {
                    prime: self.prime,
                    extension_degree: self.extension_degree,
                },
                scope,
                crate::domain::Evidence::Axiom,
            ),
        ])
    }
}

impl fmt::Display for FiniteFieldElement {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            FiniteFieldElement::Prime { residue, .. } => write!(f, "{}", residue),
            FiniteFieldElement::Extension { coeffs, .. } => {
                write!(f, "[")?;
                for (i, c) in coeffs.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", c)?;
                }
                write!(f, "]")
            }
        }
    }
}

impl Add for FiniteFieldElement {
    type Output = FiniteFieldElement;
    fn add(self, rhs: FiniteFieldElement) -> FiniteFieldElement {
        FiniteFieldElement::add(&self, &rhs)
            .expect("Add trait on extension element: use FiniteFieldDomain::add instead")
    }
}

impl Sub for FiniteFieldElement {
    type Output = FiniteFieldElement;
    fn sub(self, rhs: FiniteFieldElement) -> FiniteFieldElement {
        FiniteFieldElement::sub(&self, &rhs)
            .expect("Sub trait on extension element: use FiniteFieldDomain::sub instead")
    }
}

impl Mul for FiniteFieldElement {
    type Output = FiniteFieldElement;
    fn mul(self, rhs: FiniteFieldElement) -> FiniteFieldElement {
        self.mul_prime(&rhs)
            .expect("Mul operator on extension requires domain.mul; use FiniteFieldDomain::mul")
    }
}

impl Neg for FiniteFieldElement {
    type Output = FiniteFieldElement;
    fn neg(self) -> FiniteFieldElement {
        FiniteFieldElement::neg(&self)
    }
}

// Doc tests for FiniteFieldDomain.  These appear in the rustdoc and run
// under `cargo test --doc`.
#[cfg(test)]
mod doc_tests {
    use super::*;

    /// Build F_5 and verify that 2 + 3 = 0 in the field.
    #[test]
    fn doc_f5_add() {
        let f5 = FiniteFieldDomain::prime(5);
        let two = f5.element(&[2]);
        let three = f5.element(&[3]);
        assert_eq!(f5.add(&two, &three).unwrap(), f5.zero());
    }

    /// Build F_4 = F_2[x]/(x^2+x+1) and verify that x^2 = x + 1.
    #[test]
    fn doc_f4_x_squared() {
        let f4 = FiniteFieldDomain::extension(2, 2, vec![1, 1, 1]);
        let x = f4.element(&[0, 1]);
        let x_sq = f4.mul(&x, &x).unwrap();
        assert_eq!(x_sq, f4.element(&[1, 1]));
    }
}