Skip to main content

trident/field/
mod.rs

1//! Prime field arithmetic and universal proving primitives.
2//!
3//! This module provides field-generic math that every target warrior reuses:
4//! - `PrimeField` trait with concrete implementations (Goldilocks, BabyBear, Mersenne31)
5//! - `poseidon2` — generic Poseidon2 sponge hash over any PrimeField
6//! - `proof` — universal STARK proof estimation (padded height, FRI params, proof size)
7//!
8//! Three fields cover all 20 supported VMs:
9//! - Goldilocks (2^64 - 2^32 + 1): Triton, Miden, OpenVM, Plonky3
10//! - BabyBear (2^31 - 2^27 + 1): SP1, RISC Zero, Jolt
11//! - Mersenne31 (2^31 - 1): Plonky3, Circle STARKs
12
13pub mod babybear;
14pub mod fixed;
15pub mod goldilocks;
16pub mod mersenne31;
17pub mod poseidon2;
18pub mod proof;
19
20pub use babybear::BabyBear;
21pub use goldilocks::Goldilocks;
22pub use mersenne31::Mersenne31;
23
24/// Trait for prime field arithmetic.
25///
26/// Warriors use this for field-generic hash functions, proof estimation,
27/// and verification. Trident provides concrete implementations;
28/// warriors call them without reimplementing the math.
29pub trait PrimeField: Copy + Clone + Eq + PartialEq + Ord + PartialOrd + std::fmt::Debug {
30    /// The field modulus as u128 (fits all supported primes).
31    const MODULUS: u128;
32    /// Number of bits in the modulus.
33    const BITS: u32;
34    /// Additive identity.
35    const ZERO: Self;
36    /// Multiplicative identity.
37    const ONE: Self;
38
39    /// Construct from a u64 value (reduced mod p).
40    fn from_u64(v: u64) -> Self;
41    /// Extract the canonical u64 representative.
42    fn to_u64(self) -> u64;
43
44    /// Field addition: (a + b) mod p.
45    fn add(self, rhs: Self) -> Self;
46    /// Field subtraction: (a - b) mod p.
47    fn sub(self, rhs: Self) -> Self;
48    /// Field multiplication: (a * b) mod p.
49    fn mul(self, rhs: Self) -> Self;
50    /// Additive inverse: (-a) mod p.
51    fn neg(self) -> Self;
52
53    /// Multiplicative inverse via Fermat: a^(p-2) mod p.
54    /// Returns None for zero.
55    fn inv(self) -> Option<Self> {
56        if self == Self::ZERO {
57            return None;
58        }
59        // Default: Fermat's little theorem. Implementors may override
60        // with target-specific optimized inversion.
61        let exp = Self::MODULUS - 2;
62        Some(self.pow_u128(exp))
63    }
64
65    /// Exponentiation: a^exp mod p (square-and-multiply, u64 exponent).
66    fn pow(self, mut exp: u64) -> Self {
67        let mut base = self;
68        let mut acc = Self::ONE;
69        while exp > 0 {
70            if exp & 1 == 1 {
71                acc = acc.mul(base);
72            }
73            base = base.mul(base);
74            exp >>= 1;
75        }
76        acc
77    }
78
79    /// Exponentiation with u128 exponent (for Fermat inversion).
80    fn pow_u128(self, mut exp: u128) -> Self {
81        let mut base = self;
82        let mut acc = Self::ONE;
83        while exp > 0 {
84            if exp & 1 == 1 {
85                acc = acc.mul(base);
86            }
87            base = base.mul(base);
88            exp >>= 1;
89        }
90        acc
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    /// Generic test suite — run for every PrimeField implementation.
99    fn test_field_laws<F: PrimeField>() {
100        let zero = F::ZERO;
101        let one = F::ONE;
102        let a = F::from_u64(42);
103        let b = F::from_u64(1337);
104
105        // Additive identity
106        assert_eq!(a.add(zero), a);
107        assert_eq!(zero.add(a), a);
108
109        // Multiplicative identity
110        assert_eq!(a.mul(one), a);
111        assert_eq!(one.mul(a), a);
112
113        // Multiplicative zero
114        assert_eq!(a.mul(zero), zero);
115
116        // Commutativity
117        assert_eq!(a.add(b), b.add(a));
118        assert_eq!(a.mul(b), b.mul(a));
119
120        // Negation
121        assert_eq!(a.add(a.neg()), zero);
122        assert_eq!(zero.neg(), zero);
123
124        // Subtraction is add(neg)
125        assert_eq!(a.sub(b), a.add(b.neg()));
126
127        // Inverse
128        if let Some(inv_a) = a.inv() {
129            assert_eq!(a.mul(inv_a), one);
130        }
131        assert!(zero.inv().is_none());
132
133        // Power
134        assert_eq!(a.pow(0), one);
135        assert_eq!(a.pow(1), a);
136        assert_eq!(a.pow(3), a.mul(a).mul(a));
137
138        // (-1)^2 = 1
139        let neg_one = one.neg();
140        assert_eq!(neg_one.mul(neg_one), one);
141    }
142
143    #[test]
144    fn goldilocks_field_laws() {
145        test_field_laws::<Goldilocks>();
146    }
147
148    #[test]
149    fn babybear_field_laws() {
150        test_field_laws::<BabyBear>();
151    }
152
153    #[test]
154    fn mersenne31_field_laws() {
155        test_field_laws::<Mersenne31>();
156    }
157
158    /// Edge cases: values at modulus boundary, overflow in add, reduction on input.
159    fn test_field_edge_cases<F: PrimeField>() {
160        let zero = F::ZERO;
161        let one = F::ONE;
162        let p_minus_1 = F::from_u64((F::MODULUS - 1) as u64);
163
164        // p-1 + 1 = 0 (mod p)
165        assert_eq!(p_minus_1.add(one), zero);
166
167        // p-1 + p-1 = p-2 (mod p)
168        let p_minus_2 = F::from_u64((F::MODULUS - 2) as u64);
169        assert_eq!(p_minus_1.add(p_minus_1), p_minus_2);
170
171        // 0 - 1 = p-1
172        assert_eq!(zero.sub(one), p_minus_1);
173
174        // (p-1) * (p-1) = 1 (since (p-1) = -1 and (-1)^2 = 1)
175        assert_eq!(p_minus_1.mul(p_minus_1), one);
176
177        // from_u64 reduces values >= p
178        let p_as_u64 = F::MODULUS as u64;
179        assert_eq!(F::from_u64(p_as_u64), zero);
180        assert_eq!(F::from_u64(p_as_u64 + 1), one);
181
182        // Inverse of p-1 is p-1 (since -1 * -1 = 1)
183        assert_eq!(p_minus_1.inv(), Some(p_minus_1));
184
185        // Inverse of 1 is 1
186        assert_eq!(one.inv(), Some(one));
187
188        // pow(p-1, 2) = 1
189        assert_eq!(p_minus_1.pow(2), one);
190    }
191
192    #[test]
193    fn goldilocks_edge_cases() {
194        test_field_edge_cases::<Goldilocks>();
195
196        // Goldilocks-specific: test reduce128 with large products
197        let large = Goldilocks::from_u64(u64::MAX);
198        let result = large.mul(large);
199        // (u64::MAX mod p)^2 mod p — just verify it doesn't panic
200        assert!(result.to_u64() < goldilocks::MODULUS);
201    }
202
203    #[test]
204    fn babybear_edge_cases() {
205        test_field_edge_cases::<BabyBear>();
206    }
207
208    #[test]
209    fn mersenne31_edge_cases() {
210        test_field_edge_cases::<Mersenne31>();
211    }
212}