next_prime/
lib.rs

1/// Returns the ceiling of the square root of an unsigned integer in a
2/// timely manner, using binary search.  
3/// Time complexity: O(log(n))
4fn usqrt(n: u64) -> u64 {
5    let (mut low, mut high) = (1, n);
6    let mut mid = (low + high) / 2;
7    while low < high {
8        mid = (low + high) / 2;
9        let square = mid * mid;
10        if square == n {
11            return mid;
12        } else if square > n {
13            high = mid - 1
14        } else {
15            low = mid + 1
16        }
17    }
18    if mid * mid == n {
19        mid
20    } else {
21        high
22    }
23}
24
25#[cfg(test)]
26mod usqrt_tests {
27    use super::usqrt;
28
29    #[test]
30    fn small_perfect_squares() {
31        let inputs = (1..11).collect::<Vec<u64>>();
32        let squares = inputs.iter().map(|x| x * x).collect::<Vec<u64>>();
33        let outputs = squares.into_iter().map(usqrt).collect::<Vec<u64>>();
34        assert_eq!(inputs, outputs);
35    }
36
37    #[test]
38    fn large_perfect_squares() {
39        let inputs = (1000..10000).step_by(37).collect::<Vec<u64>>();
40        let squares = inputs.iter().map(|x| x * x).collect::<Vec<u64>>();
41        let outputs = squares.into_iter().map(usqrt).collect::<Vec<u64>>();
42        assert_eq!(inputs, outputs);
43    }
44
45    #[test]
46    fn edge_case_zero() {
47        assert_eq!(usqrt(0), 0);
48    }
49
50    #[test]
51    fn edge_case_one() {
52        assert_eq!(usqrt(1), 1);
53    }
54
55    #[test]
56    fn rounds_up_when_not_a_perfect_square() {
57        assert_eq!(usqrt(2), 2);
58    }
59
60    #[test]
61    fn large_not_perfect_square() {
62        let x = 12345;
63        assert_eq!(usqrt(x * x + 1), x + 1)
64    }
65}
66
67/// Determines whether a number is prime or not.  
68/// Time complexity: O(sqrt(n))
69fn is_prime(n: u64) -> bool {
70    if n == 2 || n == 3 {
71        return true;
72    }
73    if n % 2 == 0 || n <= 1 {
74        return false;
75    }
76    let lower = 3;
77    let upper = usqrt(n);
78    (lower..(upper + 1))
79        .step_by(2)
80        .all(|maybe_divisor| n % maybe_divisor != 0)
81}
82
83#[cfg(test)]
84mod is_prime_tests {
85    use super::is_prime;
86
87    #[test]
88    fn small_primes() {
89        let primes = vec![2, 3, 5, 7, 11, 13, 17, 19, 23];
90        assert_eq!(
91            primes
92                .clone()
93                .into_iter()
94                .map(is_prime)
95                .collect::<Vec<bool>>(),
96            vec![true; primes.len()]
97        );
98    }
99}
100
101/// Finds the next prime number >= n.  
102/// Time complexity: *expected* O(sqrt(n))
103/// # Examples
104/// ```
105/// use next_prime::next_prime;
106/// assert_eq!(next_prime(2), 2);
107/// assert_eq!(next_prime(4), 5);
108/// ```
109pub fn next_prime(mut n: u64) -> u64 {
110    if n <= 2 {
111        return 2;
112    }
113    if n % 2 == 0 {
114        n += 1;
115    }
116    while !is_prime(n) {
117        n += 2;
118    }
119    n
120}
121
122#[cfg(test)]
123mod next_prime_tests {
124    use super::next_prime;
125
126    #[test]
127    fn edge_case_two() {
128        assert_eq!(next_prime(2), 2);
129    }
130
131    #[test]
132    fn finds_small_primes() {
133        let primes = vec![5, 7, 11, 13, 17, 19, 23, 29];
134        assert_eq!(
135            primes,
136            primes
137                .iter()
138                .map(|x| next_prime(x - 1))
139                .collect::<Vec<u64>>()
140        );
141    }
142
143    #[test]
144    fn returns_argument_when_it_is_already_prime() {
145        assert_eq!(next_prime(101), 101);
146    }
147
148    #[test]
149    fn finds_a_very_large_prime() {
150        assert_eq!(next_prime(472_888_178), 472_888_217)
151    }
152}
next_prime/lib.rs

next_prime/
lib.rs
