polycore/
fp.rs

1use std::fmt;
2use std::hash::{Hash, Hasher};
3use std::ops::{Add, Div, Mul, Neg, Sub};
4use std::str::FromStr;
5
6use num_rational::BigRational;
7use num_traits::{One, Zero};
8
9use crate::field::Field;
10use crate::{crt, modp};
11
12/// A prime field of word-sized residues, for code that drops to bare `u64` arithmetic in
13/// [`modp`] on its inner loops.
14pub trait Modular: Field + Copy {
15    /// The prime, or 0 for a constant not yet bound to one.
16    fn modulus(&self) -> u64;
17    /// The stored residue; a signed integer in two's complement while unbound.
18    fn residue(&self) -> u64;
19    fn from_residue(v: u64, p: u64) -> Self;
20
21    /// The residue in `[0, p)`, binding an unbound constant to `p`.
22    fn residue_mod(&self, p: u64) -> u64 {
23        if self.modulus() == 0 {
24            i128::from(self.residue() as i64).rem_euclid(i128::from(p)) as u64
25        } else {
26            self.residue()
27        }
28    }
29}
30
31/// An element of `GF(p)` for a prime `p < 2^64` chosen at runtime.
32///
33/// `Fp::zero()` and `Fp::one()` carry no modulus (`p = 0`) and adopt the modulus of whatever they
34/// meet, so generic code that starts from `F::zero()` or `F::one()` works unchanged.
35#[derive(Clone, Copy, Debug)]
36pub struct Fp {
37    v: u64,
38    p: u64,
39}
40
41impl Fp {
42    /// `v` modulo `p`, or an unbound constant when `p = 0`.
43    pub const fn new(v: u64, p: u64) -> Self {
44        Self {
45            v: if p == 0 { v } else { v % p },
46            p,
47        }
48    }
49
50    pub fn from_i64(x: i64, p: u64) -> Self {
51        Self {
52            v: i128::from(x).rem_euclid(i128::from(p)) as u64,
53            p,
54        }
55    }
56
57    /// `q` modulo `p`, or `None` when `p` divides its denominator.
58    pub fn from_rational(q: &BigRational, p: u64) -> Option<Self> {
59        crt::reduce(q, p).map(|v| Self { v, p })
60    }
61
62    pub const fn value(self) -> u64 {
63        self.v
64    }
65
66    /// The prime, or 0 for a constant not yet bound to one.
67    pub const fn modulus(self) -> u64 {
68        self.p
69    }
70
71    /// The inverse, or `None` for a non-unit: zero, a residue sharing a factor with a composite
72    /// modulus, or an unbound constant other than `1` and `-1`.
73    pub const fn try_inverse(self) -> Option<Self> {
74        match (self.p, self.v) {
75            (0, 1 | u64::MAX) => Some(self),
76            (0, _) => None,
77            (p, v) => match modp::try_inv(v, p) {
78                Some(v) => Some(Self { v, p }),
79                None => None,
80            },
81        }
82    }
83
84    #[must_use]
85    pub fn bind(self, p: u64) -> Self {
86        if self.p == 0 && p != 0 {
87            Self::from_i64(self.v as i64, p)
88        } else {
89            self
90        }
91    }
92
93    fn unify(self, o: Self) -> (u64, u64, u64) {
94        let p = self.p.max(o.p);
95        (self.bind(p).v, o.bind(p).v, p)
96    }
97
98    #[inline]
99    fn lift(self, o: Self, wrap: fn(u64, u64) -> u64, op: fn(u64, u64, u64) -> u64) -> Self {
100        if self.p == o.p && self.p != 0 {
101            return Self {
102                v: op(self.v, o.v, self.p),
103                p: self.p,
104            };
105        }
106        self.lift_cold(o, wrap, op)
107    }
108
109    #[cold]
110    fn lift_cold(self, o: Self, wrap: fn(u64, u64) -> u64, op: fn(u64, u64, u64) -> u64) -> Self {
111        let (a, b, p) = self.unify(o);
112        let v = if p == 0 { wrap(a, b) } else { op(a, b, p) };
113        Self { v, p }
114    }
115}
116
117impl PartialEq for Fp {
118    fn eq(&self, o: &Self) -> bool {
119        if self.p == o.p {
120            return self.v == o.v;
121        }
122        let (a, b, _) = self.unify(*o);
123        a == b
124    }
125}
126
127impl Eq for Fp {}
128
129/// Hashes the symmetric residue, so equal elements of one field hash alike and an unbound
130/// constant hashes like its binding whenever it is below `p / 2` in magnitude.
131impl Hash for Fp {
132    fn hash<H: Hasher>(&self, h: &mut H) {
133        let v = if self.p == 0 {
134            i128::from(self.v as i64)
135        } else {
136            modp::symmetric(self.v, self.p)
137        };
138        v.hash(h);
139    }
140}
141
142impl From<i32> for Fp {
143    /// An unbound constant.
144    fn from(x: i32) -> Self {
145        Self {
146            v: i64::from(x) as u64,
147            p: 0,
148        }
149    }
150}
151
152impl Modular for Fp {
153    fn modulus(&self) -> u64 {
154        self.p
155    }
156    fn residue(&self) -> u64 {
157        self.v
158    }
159    fn from_residue(v: u64, p: u64) -> Self {
160        Self::new(v, p)
161    }
162}
163
164impl Add for Fp {
165    type Output = Self;
166    fn add(self, o: Self) -> Self {
167        self.lift(o, u64::wrapping_add, modp::add)
168    }
169}
170
171impl Sub for Fp {
172    type Output = Self;
173    fn sub(self, o: Self) -> Self {
174        self.lift(o, u64::wrapping_sub, modp::sub)
175    }
176}
177
178impl Mul for Fp {
179    type Output = Self;
180    fn mul(self, o: Self) -> Self {
181        self.lift(o, u64::wrapping_mul, modp::mul)
182    }
183}
184
185impl Div for Fp {
186    type Output = Self;
187    fn div(self, o: Self) -> Self {
188        assert!(!o.is_zero(), "division by zero in GF(p)");
189        let unit = |a: u64, b: u64| {
190            assert!(b == 1 || b == u64::MAX, "division of unbound constants");
191            a.wrapping_mul(b)
192        };
193        self.lift(o, unit, |a, b, p| modp::mul(a, modp::inv(b, p), p))
194    }
195}
196
197impl Neg for Fp {
198    type Output = Self;
199    fn neg(self) -> Self {
200        Self::zero() - self
201    }
202}
203
204impl Zero for Fp {
205    fn zero() -> Self {
206        Self { v: 0, p: 0 }
207    }
208    fn is_zero(&self) -> bool {
209        self.v == 0
210    }
211}
212
213impl One for Fp {
214    fn one() -> Self {
215        Self { v: 1, p: 0 }
216    }
217}
218
219impl fmt::Display for Fp {
220    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221        if self.p == 0 {
222            write!(f, "{}", self.v as i64)
223        } else {
224            write!(f, "{}", self.v)
225        }
226    }
227}
228
229/// An element of `GF(P)` for a prime `P` fixed at compile time.
230#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)]
231pub struct Gf<const P: u64>(u64);
232
233impl<const P: u64> Gf<P> {
234    pub const fn new(v: u64) -> Self {
235        Self(v % P)
236    }
237
238    pub fn from_i64(x: i64) -> Self {
239        Self(i128::from(x).rem_euclid(i128::from(P)) as u64)
240    }
241
242    pub const fn value(self) -> u64 {
243        self.0
244    }
245
246    pub const fn modulus(self) -> u64 {
247        P
248    }
249}
250
251impl<const P: u64> Modular for Gf<P> {
252    fn modulus(&self) -> u64 {
253        P
254    }
255    fn residue(&self) -> u64 {
256        self.0
257    }
258    fn from_residue(v: u64, _: u64) -> Self {
259        Self::new(v)
260    }
261}
262
263impl<const P: u64> Add for Gf<P> {
264    type Output = Self;
265    fn add(self, o: Self) -> Self {
266        Self(modp::add(self.0, o.0, P))
267    }
268}
269
270impl<const P: u64> Sub for Gf<P> {
271    type Output = Self;
272    fn sub(self, o: Self) -> Self {
273        Self(modp::sub(self.0, o.0, P))
274    }
275}
276
277impl<const P: u64> Mul for Gf<P> {
278    type Output = Self;
279    fn mul(self, o: Self) -> Self {
280        Self(modp::mul(self.0, o.0, P))
281    }
282}
283
284impl<const P: u64> Div for Gf<P> {
285    type Output = Self;
286    fn div(self, o: Self) -> Self {
287        assert!(o.0 != 0, "division by zero in GF(p)");
288        Self(modp::mul(self.0, modp::inv(o.0, P), P))
289    }
290}
291
292impl<const P: u64> Neg for Gf<P> {
293    type Output = Self;
294    fn neg(self) -> Self {
295        Self(modp::neg(self.0, P))
296    }
297}
298
299impl<const P: u64> Zero for Gf<P> {
300    fn zero() -> Self {
301        Self(0)
302    }
303    fn is_zero(&self) -> bool {
304        self.0 == 0
305    }
306}
307
308impl<const P: u64> One for Gf<P> {
309    fn one() -> Self {
310        Self(1 % P)
311    }
312}
313
314impl<const P: u64> From<u64> for Gf<P> {
315    fn from(v: u64) -> Self {
316        Self::new(v)
317    }
318}
319
320impl<const P: u64> From<u32> for Gf<P> {
321    fn from(v: u32) -> Self {
322        Self::new(v.into())
323    }
324}
325
326impl<const P: u64> From<i32> for Gf<P> {
327    fn from(x: i32) -> Self {
328        Self::from_i64(x.into())
329    }
330}
331
332#[derive(Clone, Debug, PartialEq, Eq)]
333pub struct ParseFpError;
334
335impl fmt::Display for ParseFpError {
336    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
337        f.write_str("expected an optionally signed integer")
338    }
339}
340
341impl std::error::Error for ParseFpError {}
342
343/// Reduces as it reads, so any number of digits is accepted.
344impl<const P: u64> FromStr for Gf<P> {
345    type Err = ParseFpError;
346    fn from_str(s: &str) -> Result<Self, ParseFpError> {
347        let s = s.trim();
348        let (neg, digits) = s
349            .strip_prefix('-')
350            .map_or_else(|| (false, s.strip_prefix('+').unwrap_or(s)), |d| (true, d));
351        if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) {
352            return Err(ParseFpError);
353        }
354        let v = digits.bytes().fold(Self::new(0), |acc, b| {
355            acc * Self::new(10) + Self::new(u64::from(b - b'0'))
356        });
357        Ok(if neg { -v } else { v })
358    }
359}
360
361impl<const P: u64> fmt::Display for Gf<P> {
362    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
363        write!(f, "{}", self.0)
364    }
365}
polycore/fp.rs

polycore/
fp.rs
