From ed09dd5590bf98a3e4646b90dda9d2b87b4403c2 Mon Sep 17 00:00:00 2001 From: Rorik Star Platinum Date: Sun, 14 Dec 2025 22:44:05 +0300 Subject: [PATCH] prime-factors --- rust/nth-prime/src/lib.rs | 133 ++++++++++++++-------- rust/nth-prime/src/main.rs | 82 +++++++++++++ rust/prime-factors/miller-rabin.rs | 115 +++++++++++++++++++ rust/prime-factors/src/lib.rs | 33 +++++- rust/prime-factors/tests/prime_factors.rs | 11 -- 5 files changed, 316 insertions(+), 58 deletions(-) create mode 100644 rust/nth-prime/src/main.rs create mode 100644 rust/prime-factors/miller-rabin.rs diff --git a/rust/nth-prime/src/lib.rs b/rust/nth-prime/src/lib.rs index d44f7e0..2637655 100644 --- a/rust/nth-prime/src/lib.rs +++ b/rust/nth-prime/src/lib.rs @@ -1,48 +1,91 @@ -const PRIME_LEN: usize = 10_010; -const LIMIT: usize = 106_000; - -struct PrimeCache { - primes: [u32; PRIME_LEN], -} - -impl PrimeCache { - const fn new() -> Self { - let mut primes = [0; PRIME_LEN]; - let mut sieve = [true; LIMIT]; - - sieve[0] = false; - sieve[1] = false; - - let mut i = 2; - while i * i < LIMIT { - if sieve[i] { - let mut j = i * i; - while j < LIMIT { - sieve[j] = false; - j += i; - } - } - i += 1; - } - let mut num = 0; - let mut j = 0; - while num < LIMIT && j < PRIME_LEN { - if sieve[num] { - primes[j] = num as u32; - j += 1; - } - num += 1; - } - Self { primes } - } -} - -static CACHE: PrimeCache = PrimeCache::new(); - pub fn nth(n: u32) -> u32 { - let nth = n as usize; - if nth >= PRIME_LEN { - panic!("N more than comptime was calculated"); + let mut num = 1; + for _ in 0..=n { + loop { + num += 1; + if is_prime(num as u64) { + break; + } + } } - CACHE.primes[nth] + num +} + +pub fn is_prime(n: u64) -> bool { + if n < 2 { + return false; + } + if n == 2 || n == 3 { + return true; + } + if n % 2 == 0 { + return false; + } + + let s = (n - 1).trailing_zeros(); + let d = (n - 1) >> s; + + let witnesses = get_witnesses(n); + + for &a in witnesses { + if !miller_rabin_test(a, d, n, s) { + return false; + } + } + true +} + +fn get_witnesses(n: u64) -> &'static [u64] { + const WITNESSES: &[(u64, &[u64])] = &[ + (2_046, &[2]), + (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]), + (u64::MAX, &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]), + ]; + + WITNESSES + .iter() + .find(|(limit, _)| *limit >= n) + .map(|(_, w)| *w) + .unwrap() +} + +fn miller_rabin_test(a: u64, d: u64, n: u64, s: u32) -> bool { + let mut x = mod_pow(a, d, n); + if x == 1 || x == n - 1 { + return true; + } + + for _ in 1..s { + x = mod_mul(x, x, n); + if x == n - 1 { + return true; + } + } + false +} + +#[inline] +fn mod_mul(a: u64, b: u64, m: u64) -> u64 { + ((a as u128 * b as u128) % m as u128) as u64 +} + +fn mod_pow(mut base: u64, mut exp: u64, m: u64) -> u64 { + let mut res = 1; + base %= m; + while exp > 0 { + if exp % 2 == 1 { + res = mod_mul(res, base, m); + } + base = mod_mul(base, base, m); + exp /= 2; + } + res } diff --git a/rust/nth-prime/src/main.rs b/rust/nth-prime/src/main.rs new file mode 100644 index 0000000..bd5989e --- /dev/null +++ b/rust/nth-prime/src/main.rs @@ -0,0 +1,82 @@ +fn main() { + let n: u64 = std::env::args().nth(1) + .and_then(|s| s.parse().ok()) + .unwrap_or(35); + + println!("{} -> {}", n, if is_prime(n) { "prime" } else { "composite" }); +} + +pub fn is_prime(n: u64) -> bool { + if n < 2 { return false; } + if n == 2 || n == 3 { return true; } + if n % 2 == 0 { return false; } + + let s = (n - 1).trailing_zeros(); + let d = (n - 1) >> s; + + let witnesses = get_witnesses(n); + + for &a in witnesses { + if !miller_rabin_test(a, d, n, s) { + return false; + } + } + true +} + +// O(1) lookup based on input size +fn get_witnesses(n: u64) -> &'static [u64] { + const WITNESSES: &[(u64, &[u64])] = &[ + (2_046, &[2]), + (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]), + (u64::MAX, &[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37]), + ]; + + WITNESSES.iter() + .find(|(limit, _)| *limit >= n) + .map(|(_, w)| *w) + .unwrap() +} + +fn miller_rabin_test(a: u64, d: u64, n: u64, s: u32) -> bool { + let mut x = mod_pow(a, d, n); + if x == 1 || x == n - 1 { + return true; + } + + for _ in 1..s { + x = mod_mul(x, x, n); + if x == n - 1 { + return true; + } + } + false +} + +// === Math Core === + +#[inline] +fn mod_mul(a: u64, b: u64, m: u64) -> u64 { + ((a as u128 * b as u128) % m as u128) as u64 +} + +fn mod_pow(mut base: u64, mut exp: u64, m: u64) -> u64 { + let mut res = 1; + base %= m; + while exp > 0 { + if exp % 2 == 1 { + res = mod_mul(res, base, m); + } + base = mod_mul(base, base, m); + exp /= 2; + } + res +} diff --git a/rust/prime-factors/miller-rabin.rs b/rust/prime-factors/miller-rabin.rs new file mode 100644 index 0000000..d6aa259 --- /dev/null +++ b/rust/prime-factors/miller-rabin.rs @@ -0,0 +1,115 @@ +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 +} diff --git a/rust/prime-factors/src/lib.rs b/rust/prime-factors/src/lib.rs index 30d06a2..bba8e1e 100644 --- a/rust/prime-factors/src/lib.rs +++ b/rust/prime-factors/src/lib.rs @@ -1,3 +1,32 @@ -pub fn factors(n: u64) -> Vec { - todo!("This should calculate the prime factors of {n}") +pub trait PrimeCollector { + fn get_counted(&mut self, prime: u64) -> Vec; +} + +impl PrimeCollector for u64 { + fn get_counted(&mut self, prime: u64) -> Vec { + let mut res: Vec<_> = Vec::new(); + while 0 == *self % prime { + *self /= prime; + res.push(prime); + } + res + } +} + +pub fn factors(n: u64) -> Vec { + let mut res: Vec<_> = Vec::new(); + + let mut n_remained = n; + let mut i = 2; + while i * i <= n_remained { + if n_remained % i == 0 { + res.extend(n_remained.get_counted(i)); + } + i += 1; + } + + if n_remained > 1 { + res.push(n_remained); + } + res } diff --git a/rust/prime-factors/tests/prime_factors.rs b/rust/prime-factors/tests/prime_factors.rs index 119c9c8..b6e941a 100644 --- a/rust/prime-factors/tests/prime_factors.rs +++ b/rust/prime-factors/tests/prime_factors.rs @@ -8,7 +8,6 @@ fn no_factors() { } #[test] -#[ignore] fn prime_number() { let factors = factors(2); let expected = [2]; @@ -16,7 +15,6 @@ fn prime_number() { } #[test] -#[ignore] fn another_prime_number() { let factors = factors(3); let expected = [3]; @@ -24,7 +22,6 @@ fn another_prime_number() { } #[test] -#[ignore] fn square_of_a_prime() { let factors = factors(9); let expected = [3, 3]; @@ -32,7 +29,6 @@ fn square_of_a_prime() { } #[test] -#[ignore] fn product_of_first_prime() { let factors = factors(4); let expected = [2, 2]; @@ -40,7 +36,6 @@ fn product_of_first_prime() { } #[test] -#[ignore] fn cube_of_a_prime() { let factors = factors(8); let expected = [2, 2, 2]; @@ -48,7 +43,6 @@ fn cube_of_a_prime() { } #[test] -#[ignore] fn product_of_second_prime() { let factors = factors(27); let expected = [3, 3, 3]; @@ -56,7 +50,6 @@ fn product_of_second_prime() { } #[test] -#[ignore] fn product_of_third_prime() { let factors = factors(625); let expected = [5, 5, 5, 5]; @@ -64,7 +57,6 @@ fn product_of_third_prime() { } #[test] -#[ignore] fn product_of_first_and_second_prime() { let factors = factors(6); let expected = [2, 3]; @@ -72,7 +64,6 @@ fn product_of_first_and_second_prime() { } #[test] -#[ignore] fn product_of_primes_and_non_primes() { let factors = factors(12); let expected = [2, 2, 3]; @@ -80,7 +71,6 @@ fn product_of_primes_and_non_primes() { } #[test] -#[ignore] fn product_of_primes() { let factors = factors(901_255); let expected = [5, 17, 23, 461]; @@ -88,7 +78,6 @@ fn product_of_primes() { } #[test] -#[ignore] fn factors_include_a_large_prime() { let factors = factors(93_819_012_551); let expected = [11, 9_539, 894_119];