1//! Prime-field arithmetic for p < 2^62 in Montgomery form (R = 2^64).2//!3//! All field elements handled by the walks are kept in Montgomery representation4//! `a·R mod p`; conversion happens only at the boundaries (CLI input/output, JSON) and5//! where a canonical (standard-form) value is needed for an ordering decision.6//! Products of two residues fit in a u128; the Montgomery reduction needs one7//! conditional subtraction because p < 2^62 guarantees (a·b + m·p)/R < 2p.8//! Exact integer arithmetic only (CONVENTIONS.md §1).910#[derive(Clone, Copy, Debug)]11pub struct Field {12 pub p: u64,13 /// -p^{-1} mod 2^6414 pub pinv: u64,15 /// R mod p (Montgomery form of 1)16 pub one: u64,17 /// R^2 mod p18 pub r2: u64,19 /// R^3 mod p (used to convert a plain inverse back to Montgomery form)20 pub r3: u64,21}2223impl Field {24 pub fn new(p: u64) -> Field {25 assert!(p > 3 && p & 1 == 1, "p must be an odd prime > 3");26 assert!(p < (1u64 << 62), "p must be < 2^62");27 // Newton iteration for p^{-1} mod 2^6428 let mut inv: u64 = 1;29 for _ in 0..6 {30 inv = inv.wrapping_mul(2u64.wrapping_sub(p.wrapping_mul(inv)));31 }32 debug_assert_eq!(p.wrapping_mul(inv), 1);33 let pinv = inv.wrapping_neg();34 let r = ((1u128 << 64) % p as u128) as u64;35 let r2 = ((r as u128 * r as u128) % p as u128) as u64;36 let r3 = ((r2 as u128 * r as u128) % p as u128) as u64;37 Field { p, pinv, one: r, r2, r3 }38 }3940 /// Montgomery reduction of t < p·2^64: returns t·R^{-1} mod p.41 #[inline(always)]42 pub fn redc(&self, t: u128) -> u64 {43 let m = (t as u64).wrapping_mul(self.pinv);44 let u = ((t + (m as u128) * (self.p as u128)) >> 64) as u64;45 if u >= self.p {46 u - self.p47 } else {48 u49 }50 }5152 #[inline(always)]53 pub fn mul(&self, a: u64, b: u64) -> u64 {54 self.redc(a as u128 * b as u128)55 }5657 #[inline(always)]58 pub fn sqr(&self, a: u64) -> u64 {59 self.redc(a as u128 * a as u128)60 }6162 #[inline(always)]63 pub fn add(&self, a: u64, b: u64) -> u64 {64 let s = a + b; // < 2^63, no overflow65 if s >= self.p {66 s - self.p67 } else {68 s69 }70 }7172 #[inline(always)]73 pub fn sub(&self, a: u64, b: u64) -> u64 {74 if a >= b {75 a - b76 } else {77 a + self.p - b78 }79 }8081 #[inline(always)]82 pub fn neg(&self, a: u64) -> u64 {83 if a == 0 {84 085 } else {86 self.p - a87 }88 }8990 #[inline(always)]91 pub fn dbl(&self, a: u64) -> u64 {92 self.add(a, a)93 }9495 /// Standard integer a (0 ≤ a < p) → Montgomery form.96 #[inline(always)]97 pub fn to_mont(&self, a: u64) -> u64 {98 debug_assert!(a < self.p);99 self.mul(a, self.r2)100 }101102 /// Montgomery form → standard integer.103 #[inline(always)]104 pub fn from_mont(&self, a: u64) -> u64 {105 self.redc(a as u128)106 }107108 /// Inverse of a standard-form element by the extended Euclidean algorithm (a ≠ 0).109 pub fn inv_plain(&self, a: u64) -> u64 {110 debug_assert!(a != 0 && a < self.p);111 let (mut r0, mut r1) = (self.p as i64, a as i64);112 let (mut t0, mut t1) = (0i64, 1i64);113 while r1 != 0 {114 let q = r0 / r1;115 let r2 = r0 - q * r1;116 r0 = r1;117 r1 = r2;118 let t2 = t0 - q * t1;119 t0 = t1;120 t1 = t2;121 }122 debug_assert_eq!(r0, 1, "inverse of non-unit");123 if t0 < 0 {124 (t0 + self.p as i64) as u64125 } else {126 t0 as u64127 }128 }129130 /// Inverse of a Montgomery-form element, result in Montgomery form.131 /// inv_plain(aR) = a^{-1} R^{-1}; multiplying by R^3 (one Montgomery product) gives a^{-1} R.132 #[inline]133 pub fn inv(&self, a: u64) -> u64 {134 let t = self.inv_plain(a);135 self.mul(t, self.r3)136 }137138 /// Montgomery's simultaneous inversion: replaces every element of `xs` (all non-zero,139 /// Montgomery form) by its inverse, using 3(k−1) multiplications and a single inversion.140 /// `scratch` must have length ≥ xs.len(). Returns (muls, inversions) performed.141 pub fn batch_inv(&self, xs: &mut [u64], scratch: &mut [u64]) -> (u64, u64) {142 let k = xs.len();143 if k == 0 {144 return (0, 0);145 }146 scratch[0] = xs[0];147 for i in 1..k {148 scratch[i] = self.mul(scratch[i - 1], xs[i]);149 }150 let mut acc = self.inv(scratch[k - 1]);151 for i in (1..k).rev() {152 let inv_i = self.mul(acc, scratch[i - 1]);153 acc = self.mul(acc, xs[i]);154 xs[i] = inv_i;155 }156 xs[0] = acc;157 ((3 * (k as u64 - 1)) + 1, 1)158 }159160 pub fn pow_plain(&self, a: u64, mut e: u64) -> u64 {161 let mut base = self.to_mont(a);162 let mut acc = self.one;163 while e > 0 {164 if e & 1 == 1 {165 acc = self.mul(acc, base);166 }167 base = self.sqr(base);168 e >>= 1;169 }170 self.from_mont(acc)171 }172173 /// Legendre symbol as 0, 1 or p−1 (standard form input).174 pub fn legendre(&self, a: u64) -> u64 {175 self.pow_plain(a % self.p, (self.p - 1) / 2)176 }177178 /// Tonelli–Shanks square root of a standard-form element; None if non-residue.179 pub fn sqrt_plain(&self, a: u64) -> Option<u64> {180 let p = self.p;181 let a = a % p;182 if a == 0 {183 return Some(0);184 }185 if self.legendre(a) != 1 {186 return None;187 }188 if p % 4 == 3 {189 return Some(self.pow_plain(a, (p + 1) / 4));190 }191 let (mut q, mut s) = (p - 1, 0u32);192 while q % 2 == 0 {193 q /= 2;194 s += 1;195 }196 let mut z = 2u64;197 while self.legendre(z) != p - 1 {198 z += 1;199 }200 let mut m = s;201 let mut c = self.to_mont(self.pow_plain(z, q));202 let mut t = self.to_mont(self.pow_plain(a, q));203 let mut r = self.to_mont(self.pow_plain(a, (q + 1) / 2));204 while t != self.one {205 let mut i = 0u32;206 let mut t2 = t;207 while t2 != self.one {208 t2 = self.sqr(t2);209 i += 1;210 }211 let mut b = c;212 for _ in 0..(m - i - 1) {213 b = self.sqr(b);214 }215 m = i;216 c = self.sqr(b);217 t = self.mul(t, c);218 r = self.mul(r, b);219 }220 Some(self.from_mont(r))221 }222}223224/// Arithmetic modulo the (prime) group order n < 2^62 — used for the scalar bookkeeping.225#[derive(Clone, Copy, Debug)]226pub struct Modn {227 pub n: u64,228}229230impl Modn {231 pub fn new(n: u64) -> Modn {232 assert!(n > 1 && n < (1u64 << 62));233 Modn { n }234 }235 #[inline(always)]236 pub fn add(&self, a: u64, b: u64) -> u64 {237 let s = a + b;238 if s >= self.n {239 s - self.n240 } else {241 s242 }243 }244 #[inline(always)]245 pub fn sub(&self, a: u64, b: u64) -> u64 {246 if a >= b {247 a - b248 } else {249 a + self.n - b250 }251 }252 #[inline(always)]253 pub fn neg(&self, a: u64) -> u64 {254 if a == 0 {255 0256 } else {257 self.n - a258 }259 }260 #[inline(always)]261 pub fn mul(&self, a: u64, b: u64) -> u64 {262 ((a as u128 * b as u128) % self.n as u128) as u64263 }264 #[inline(always)]265 pub fn dbl(&self, a: u64) -> u64 {266 self.add(a, a)267 }268 pub fn inv(&self, a: u64) -> Option<u64> {269 if a == 0 {270 return None;271 }272 let (mut r0, mut r1) = (self.n as i64, a as i64);273 let (mut t0, mut t1) = (0i64, 1i64);274 while r1 != 0 {275 let q = r0 / r1;276 let r2 = r0 - q * r1;277 r0 = r1;278 r1 = r2;279 let t2 = t0 - q * t1;280 t0 = t1;281 t1 = t2;282 }283 if r0 != 1 {284 return None;285 }286 Some(if t0 < 0 { (t0 + self.n as i64) as u64 } else { t0 as u64 })287 }288 /// Solve a·d ≡ b (mod n); None if a ≡ 0.289 pub fn solve(&self, a: u64, b: u64) -> Option<u64> {290 self.inv(a).map(|ai| self.mul(ai, b))291 }292 pub fn pow(&self, a: u64, mut e: u64) -> u64 {293 let mut base = a % self.n;294 let mut acc = 1u64;295 while e > 0 {296 if e & 1 == 1 {297 acc = self.mul(acc, base);298 }299 base = self.mul(base, base);300 e >>= 1;301 }302 acc303 }304}305306#[cfg(test)]307mod tests {308 use super::*;309310 #[test]311 fn montgomery_roundtrip_and_inverse() {312 for &p in &[1_000_003u64, 4611686018427387847u64 /* 2^62-57 */, 2305843009213693951u64] {313 let f = Field::new(p);314 let mut x = 123456789u64 % p;315 for _ in 0..1000 {316 x = (x as u128 * 6364136223846793005u128 % p as u128) as u64 + 1;317 x %= p;318 if x == 0 {319 continue;320 }321 let xm = f.to_mont(x);322 assert_eq!(f.from_mont(xm), x);323 let y = f.inv(xm);324 assert_eq!(f.from_mont(f.mul(xm, y)), 1);325 let prod = f.from_mont(f.mul(xm, xm));326 assert_eq!(prod, ((x as u128 * x as u128) % p as u128) as u64);327 }328 }329 }330331 #[test]332 fn batch_inverse_matches_single() {333 let p = 4611686018427387847u64;334 let f = Field::new(p);335 let mut xs: Vec<u64> = (1..300u64).map(|i| f.to_mont(i * 987654321 % p)).collect();336 let expect: Vec<u64> = xs.iter().map(|&x| f.inv(x)).collect();337 let mut scratch = vec![0u64; xs.len()];338 f.batch_inv(&mut xs, &mut scratch);339 assert_eq!(xs, expect);340 }341342 #[test]343 fn sqrt_works() {344 for &p in &[1_000_003u64, 1_000_033u64, 4611686018427387847u64] {345 let f = Field::new(p);346 for a in 2..200u64 {347 if let Some(r) = f.sqrt_plain(a) {348 assert_eq!(((r as u128 * r as u128) % p as u128) as u64, a % p);349 }350 }351 }352 }353}354