1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

/// Finds the lowest prime number which is greater than the provided number
/// `n`.
pub fn find_next_prime(mut n: f64) -> f64 {
    n = n.round();

    // Handle small numbers separately
    if n <= 2.0 {
        return 2.0;
    } else if n <= 3.0 {
        return 3.0;
    }

    // All prime numbers greater than 3 are of the form 6k + 1 or 6k + 5 (or
    // 6k - 1).
    // That's because:
    //
    // 6k is divisible by 2 and 3.
    // 6k + 2 = 2(3k + 1) is divisible by 2.
    // 6k + 3 = 3(2k + 1) is divisible by 3.
    // 6k + 4 = 2(3k + 2) is divisible by 2.
    //
    // This leaves only 6k + 1 and 6k + 5 as candidates.

    // Ensure the candidate is of the form 6k - 1 or 6k + 1.
    let remainder = n % 6.0;
    if remainder != 0.0 {
        n = n + 6.0 - remainder;

        let candidate = n - 1.0;
        if is_prime(candidate) {
            return candidate;
        }
    }

    loop {
        let candidate = n + 1.0;
        if is_prime(candidate) {
            return candidate;
        }
        let candidate = n + 5.0;
        if is_prime(candidate) {
            return candidate;
        }

        n += 6.0;
    }
}

pub fn find_next_prime_with_load_factor(n: f64, load_factor: f64) -> f64 {
    let minimum = n / load_factor;
    find_next_prime(minimum)
}

/// Checks whether the provided number `n` is a prime number.
pub fn is_prime(n: f64) -> bool {
    if n <= 1.0 {
        return false;
    }
    if n <= 3.0 {
        return true;
    }
    if n % 2.0 == 0.0 || n % 3.0 == 0.0 {
        return false;
    }
    let mut i = 5.0;
    while i * i <= n {
        if n % i == 0.0 || n % (i + 2.0) == 0.0 {
            return false;
        }
        i += 6.0;
    }
    true
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_find_next_prime() {
        assert_eq!(find_next_prime(0.0), 2.0);
        assert_eq!(find_next_prime(2.0), 2.0);
        assert_eq!(find_next_prime(3.0), 3.0);
        assert_eq!(find_next_prime(4.0), 5.0);

        assert_eq!(find_next_prime(10.0), 11.0);
        assert_eq!(find_next_prime(28.0), 29.0);

        assert_eq!(find_next_prime(102.0), 103.0);
        assert_eq!(find_next_prime(105.0), 107.0);

        assert_eq!(find_next_prime(100.0), 101.0);
        assert_eq!(find_next_prime(1000.0), 1009.0);

        assert_eq!(find_next_prime(4800.0), 4801.0);
        assert_eq!(find_next_prime(5000.0), 5003.0);
        assert_eq!(find_next_prime(6000.0), 6007.0);
        assert_eq!(find_next_prime(6850.0), 6857.0);

        assert_eq!(find_next_prime(7900.0), 7901.0);
        assert_eq!(find_next_prime(7907.0), 7907.0);
    }

    #[test]
    fn test_find_next_prime_with_load_factor() {
        assert_eq!(find_next_prime_with_load_factor(4800.0, 0.5), 9601.0);
        assert_eq!(find_next_prime_with_load_factor(4800.0, 0.7), 6857.0);
    }

    #[test]
    fn test_is_prime() {
        assert_eq!(is_prime(1.0), false);
        assert_eq!(is_prime(2.0), true);
        assert_eq!(is_prime(3.0), true);
        assert_eq!(is_prime(4.0), false);
    }
}