light_utils/
prime.rs

1/// Finds the lowest prime number which is greater than the provided number
2/// `n`.
3pub fn find_next_prime(mut n: u32) -> u32 {
4    // Handle small numbers separately
5    if n <= 2 {
6        return 2;
7    } else if n <= 3 {
8        return 3;
9    }
10
11    // All prime numbers greater than 3 are of the form 6k + 1 or 6k + 5 (or
12    // 6k - 1).
13    // That's because:
14    //
15    // 6k is divisible by 2 and 3.
16    // 6k + 2 = 2(3k + 1) is divisible by 2.
17    // 6k + 3 = 3(2k + 1) is divisible by 3.
18    // 6k + 4 = 2(3k + 2) is divisible by 2.
19    //
20    // This leaves only 6k + 1 and 6k + 5 as candidates.
21
22    // Ensure the candidate is of the form 6k - 1 or 6k + 1.
23    let remainder = n % 6;
24    if remainder != 0 {
25        // Check if `n` already satisfies the pattern and is prime.
26        if remainder == 5 && is_prime(n) {
27            return n;
28        }
29        if remainder == 1 && is_prime(n) {
30            return n;
31        }
32
33        // Add `6 - remainder` to `n`, to it satisfies the `6k` pattern.
34        n = n + 6 - remainder;
35        // Check if `6k - 1` candidate is prime.
36        let candidate = n - 1;
37        if is_prime(candidate) {
38            return candidate;
39        }
40    }
41
42    // Consequently add `6`, keep checking `6k + 1` and `6k + 5` candidates.
43    loop {
44        let candidate = n + 1;
45        if is_prime(candidate) {
46            return candidate;
47        }
48        let candidate = n + 5;
49        if is_prime(candidate) {
50            return candidate;
51        }
52
53        n += 6;
54    }
55}
56
57pub fn find_next_prime_with_load_factor(n: u32, load_factor: f64) -> u32 {
58    // SAFETY: These type coercions should not cause any issues.
59    //
60    // * `f64` can precisely represent all integer values up to 2^53, which is
61    //   more than `u32::MAX`. `u64` and `usize` would be too large though.
62    // * We want to return and find an integer (prime number), so coercing `f64`
63    //   back to `u32` is intentional here.
64    let minimum = n as f64 / load_factor;
65    find_next_prime(minimum as u32)
66}
67
68/// Checks whether the provided number `n` is a prime number.
69pub fn is_prime(n: u32) -> bool {
70    if n <= 1 {
71        return false;
72    }
73    if n <= 3 {
74        return true;
75    }
76    if n % 2 == 0 || n % 3 == 0 {
77        return false;
78    }
79    let mut i = 5;
80    while i * i <= n {
81        if n % i == 0 || n % (i + 2) == 0 {
82            return false;
83        }
84        i += 6;
85    }
86    true
87}
88
89#[cfg(test)]
90mod test {
91    use super::*;
92
93    #[test]
94    fn test_find_next_prime() {
95        assert_eq!(find_next_prime(0), 2);
96        assert_eq!(find_next_prime(2), 2);
97        assert_eq!(find_next_prime(3), 3);
98        assert_eq!(find_next_prime(4), 5);
99
100        assert_eq!(find_next_prime(10), 11);
101        assert_eq!(find_next_prime(17), 17);
102        assert_eq!(find_next_prime(19), 19);
103        assert_eq!(find_next_prime(28), 29);
104
105        assert_eq!(find_next_prime(100), 101);
106        assert_eq!(find_next_prime(102), 103);
107        assert_eq!(find_next_prime(105), 107);
108
109        assert_eq!(find_next_prime(1000), 1009);
110        assert_eq!(find_next_prime(2000), 2003);
111        assert_eq!(find_next_prime(3000), 3001);
112        assert_eq!(find_next_prime(4000), 4001);
113
114        assert_eq!(find_next_prime(4800), 4801);
115        assert_eq!(find_next_prime(5000), 5003);
116        assert_eq!(find_next_prime(6000), 6007);
117        assert_eq!(find_next_prime(6850), 6857);
118
119        assert_eq!(find_next_prime(7000), 7001);
120        assert_eq!(find_next_prime(7900), 7901);
121        assert_eq!(find_next_prime(7907), 7907);
122    }
123
124    #[test]
125    fn test_find_next_prime_with_load_factor() {
126        assert_eq!(find_next_prime_with_load_factor(4800, 0.5), 9601);
127        assert_eq!(find_next_prime_with_load_factor(4800, 0.7), 6857);
128    }
129
130    #[test]
131    fn test_is_prime() {
132        assert_eq!(is_prime(1), false);
133        assert_eq!(is_prime(2), true);
134        assert_eq!(is_prime(3), true);
135        assert_eq!(is_prime(4), false);
136        assert_eq!(is_prime(17), true);
137        assert_eq!(is_prime(19), true);
138    }
139}
light_utils/prime.rs

light_utils/
prime.rs
