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