Skip to main content

lib_modulo/
prime.rs

1use crate::{Modulus32, Modulus64};
2
3/// Performs deterministic Miller-Rabin primality test.
4///
5/// # Time complexity
6///
7/// *O*(log *x*)
8#[must_use]
9#[allow(clippy::unreadable_literal)]
10#[allow(clippy::cast_possible_truncation)]
11pub fn primality_test(x: u64) -> bool {
12    if x < 64 {
13        return (PRIME_LT_64 >> x) & 1 == 1;
14    } else if (COPRIME_2_3_5 >> (x % 30)) & 1 == 0 || x % 7 == 0 {
15        return false;
16    }
17
18    let (s, d) = {
19        let x = x - 1;
20        let s = x.trailing_zeros();
21        (s - 1, x >> s)
22    };
23
24    macro_rules! primality_test_impl {
25        ($modulus:expr, $witness:expr, $d:expr, $s:expr) => {
26            let modulus = $modulus;
27            let witness = $witness;
28            let (d, s) = ($d, $s);
29
30            let one = modulus.residue(1);
31            let minus_one = -one;
32
33            let mut i = 0;
34            'test: while i < witness.len() {
35                let mut mint = modulus.residue(witness[i]);
36                i += 1;
37
38                if mint.is_zero() {
39                    continue;
40                }
41
42                mint = mint.pow(d);
43                if mint == one || mint == minus_one {
44                    continue;
45                }
46
47                let mut s = s;
48                while s > 0 {
49                    s -= 1;
50
51                    mint = mint * mint;
52                    if mint == minus_one {
53                        continue 'test;
54                    }
55                }
56
57                return false;
58            }
59        };
60    }
61
62    // witnesses from <https://miller-rabin.appspot.com/>
63    if x <= Modulus32::MAX as u64 {
64        primality_test_impl!(Modulus32::new(x as u32), [2, 7, 61], d as u32, s);
65    } else {
66        let witness = if x < 350_269_456_337 {
67            static SET3: [u64; 3] = [
68                4230279247111683200,
69                14694767155120705706,
70                16641139526367750375,
71            ];
72            SET3.as_slice()
73        } else if x < 7_999_252_175_582_851 {
74            static SET5: [u64; 5] = [
75                2,
76                4130806001517,
77                149795463772692060,
78                186635894390467037,
79                3967304179347715805,
80            ];
81            SET5.as_slice()
82        } else {
83            static SET7: [u64; 7] = [2, 325, 9375, 28178, 450775, 9780504, 1795265022];
84            SET7.as_slice()
85        };
86
87        primality_test_impl!(Modulus64::new(x), witness, d, s);
88    }
89
90    true
91}
92
93/// small prime numbers less than 64
94///
95/// `(SELF >> x) & 1 == 1` iff `x` is prime
96const PRIME_LT_64: u64 = {
97    let mut test = u64::MAX << 2;
98    let mut x = 2;
99    while x < 64 {
100        if (test >> x) & 1 == 1 {
101            let mut y = x * x;
102            while y < 64 {
103                test &= !(1 << y);
104                y += x;
105            }
106        }
107
108        x += 1;
109    }
110    test
111};
112
113/// multiples of 2, 3 or 5.
114///
115/// `(SELF >> x % 30) & 1 == 1` iff `x` is coprime to 2, 3, and 5
116const COPRIME_2_3_5: u32 = {
117    let mut table = 0;
118
119    let mut n = 0;
120    while n < 30 {
121        table |= if n % 2 == 0 || n % 3 == 0 || n % 5 == 0 {
122            0 // composite
123        } else {
124            1 // may be prime
125        } << n;
126        n += 1;
127    }
128
129    table
130};
131
132#[cfg(test)]
133mod tests {
134    use proptest::prelude::*;
135
136    use super::*;
137
138    const fn primality_test_naive(x: u64) -> bool {
139        if x < 2 {
140            return false;
141        }
142
143        let mut d = 1;
144        while d < x.isqrt() {
145            d += 1;
146
147            if x % d == 0 {
148                return false;
149            }
150        }
151
152        true
153    }
154
155    #[test]
156    fn small() {
157        for x in 0..1 << 15 {
158            assert_eq!(primality_test(x), primality_test_naive(x), "{x}")
159        }
160    }
161
162    proptest! {
163        #![proptest_config(ProptestConfig::with_cases(1 << 15))]
164        #[test]
165        fn random1(x in 1u64 << 15..1 << 20) {
166            assert_eq!(primality_test(x), primality_test_naive(x), "{x}")
167        }
168    }
169
170    proptest! {
171        #![proptest_config(ProptestConfig::with_cases(1 << 15))]
172        #[test]
173        fn random2(x in 1u64 << 20..1 << 25) {
174            assert_eq!(primality_test(x), primality_test_naive(x), "{x}")
175        }
176    }
177
178    proptest! {
179        #![proptest_config(ProptestConfig::with_cases(1 << 13))]
180        #[test]
181        fn random3(x in 1u64 << 25..1 << 30) {
182            assert_eq!(primality_test(x), primality_test_naive(x), "{x}")
183        }
184    }
185
186    proptest! {
187        #![proptest_config(ProptestConfig::with_cases(1 << 10))]
188        #[test]
189        fn random4(x in 1u64 << 30..1 << 40) {
190            assert_eq!(primality_test(x), primality_test_naive(x), "{x}")
191        }
192    }
193
194    #[test]
195    fn mersenne() {
196        for d in 0..=64 {
197            // Mersenne prime in `u64`
198            let is_prime = [2, 3, 5, 7, 13, 17, 19, 31, 61].contains(&d);
199            let x = (1u128 << d) - 1;
200
201            assert_eq!(primality_test(x as u64), is_prime)
202        }
203    }
204
205    proptest! {
206        #![proptest_config(ProptestConfig::with_cases(1 << 18))]
207        #[test]
208        fn composite(x: u32) {
209            assert!(!primality_test(x as u64 * (x as u64 + 1)));
210        }
211    }
212}