115 lines
2.8 KiB
Rust
115 lines
2.8 KiB
Rust
pub fn nth(n: u32) -> u32 {
|
|
let mut num = 1;
|
|
for _ in 0..=n {
|
|
loop {
|
|
num += 1;
|
|
if miller_rabin(num as u64) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
num
|
|
}
|
|
|
|
fn miller_rabin(n: u64) -> bool {
|
|
const HINT: &[u64] = &[2];
|
|
|
|
// we have a strict upper bound, so we can just use the witness
|
|
// table of Pomerance, Selfridge & Wagstaff and Jeaschke to be as
|
|
// efficient as possible, without having to fall back to
|
|
// randomness. Additional limits from Feitsma and Galway complete
|
|
// the entire range of `u64`. See also:
|
|
// https://en.wikipedia.org/wiki/Miller%E2%80%93Rabin_primality_test#Testing_against_small_sets_of_bases
|
|
const WITNESSES: &[(u64, &[u64])] = &[
|
|
(2_046, HINT),
|
|
(1_373_652, &[2, 3]),
|
|
(9_080_190, &[31, 73]),
|
|
(25_326_000, &[2, 3, 5]),
|
|
(4_759_123_140, &[2, 7, 61]),
|
|
(1_112_004_669_632, &[2, 13, 23, 1662803]),
|
|
(2_152_302_898_746, &[2, 3, 5, 7, 11]),
|
|
(3_474_749_660_382, &[2, 3, 5, 7, 11, 13]),
|
|
(341_550_071_728_320, &[2, 3, 5, 7, 11, 13, 17]),
|
|
(3_825_123_056_546_413_050, &[2, 3, 5, 7, 11, 13, 17, 19, 23]),
|
|
(std::u64::MAX, &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]),
|
|
];
|
|
|
|
if n % 2 == 0 {
|
|
return n == 2;
|
|
}
|
|
if n == 1 {
|
|
return false;
|
|
}
|
|
|
|
let mut d = n - 1;
|
|
let mut s = 0;
|
|
while d % 2 == 0 {
|
|
d /= 2;
|
|
s += 1
|
|
}
|
|
|
|
let witnesses = WITNESSES
|
|
.iter()
|
|
.find(|&&(hi, _)| hi >= n)
|
|
.map(|&(_, wtnss)| wtnss)
|
|
.unwrap();
|
|
'next_witness: for &a in witnesses.iter() {
|
|
let mut power = mod_exp(a, d, n);
|
|
assert!(power < n);
|
|
if power == 1 || power == n - 1 {
|
|
continue 'next_witness;
|
|
}
|
|
|
|
for _r in 0..s {
|
|
power = mod_sqr(power, n);
|
|
assert!(power < n);
|
|
if power == 1 {
|
|
return false;
|
|
}
|
|
if power == n - 1 {
|
|
continue 'next_witness;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
fn mod_mul_(a: u64, b: u64, m: u64) -> u64 {
|
|
(u128::from(a) * u128::from(b) % u128::from(m)) as u64
|
|
}
|
|
|
|
fn mod_mul(a: u64, b: u64, m: u64) -> u64 {
|
|
match a.checked_mul(b) {
|
|
Some(r) => {
|
|
if r >= m {
|
|
r % m
|
|
} else {
|
|
r
|
|
}
|
|
}
|
|
None => mod_mul_(a, b, m),
|
|
}
|
|
}
|
|
|
|
fn mod_sqr(a: u64, m: u64) -> u64 {
|
|
if a < (1 << 32) {
|
|
let r = a * a;
|
|
if r >= m { r % m } else { r }
|
|
} else {
|
|
mod_mul_(a, a, m)
|
|
}
|
|
}
|
|
|
|
fn mod_exp(mut x: u64, mut d: u64, n: u64) -> u64 {
|
|
let mut ret: u64 = 1;
|
|
while d != 0 {
|
|
if d % 2 == 1 {
|
|
ret = mod_mul(ret, x, n)
|
|
}
|
|
d /= 2;
|
|
x = mod_sqr(x, n);
|
|
}
|
|
ret
|
|
}
|