const_primes/
integer_math.rs

1//! This module contains const math operations on integers that are used by the other
2//! functions in the crate.
3
4/// Returns the largest integer smaller than or equal to `√n`.
5///
6/// # Examples
7///
8/// Basic usage:
9///
10/// ```
11/// # use const_primes::isqrt;
12/// const ISQRT25: u64 = isqrt(25);
13/// const ISQRT35: u64 = isqrt(35);
14/// const ISQRT36: u64 = isqrt(36);
15///
16/// assert_eq!(ISQRT25, 5);
17/// assert_eq!(ISQRT35, 5);
18/// assert_eq!(ISQRT36, 6);
19/// ```
20#[must_use]
21pub const fn isqrt(n: u64) -> u64 {
22    if n <= 1 {
23        n
24    } else {
25        let mut x0 = u64::pow(2, n.ilog2() / 2 + 1);
26        let mut x1 = (x0 + n / x0) / 2;
27        while x1 < x0 {
28            x0 = x1;
29            x1 = (x0 + n / x0) / 2;
30        }
31        x0
32    }
33}
34
35#[cfg(not(feature = "fast_test"))]
36/// Calculates (`base` ^ `exp`) mod `modulo` without overflow.
37#[must_use]
38pub const fn mod_pow(mut base: u64, mut exp: u64, modulo: u64) -> u64 {
39    let mut res = 1;
40
41    base %= modulo;
42
43    while exp > 0 {
44        if exp % 2 == 1 {
45            res = mod_mul(res, base, modulo);
46        }
47        base = mod_mul(base, base, modulo);
48        exp >>= 1;
49    }
50
51    res
52}
53
54#[cfg(not(feature = "fast_test"))]
55/// Calculates (`a` * `b`) mod `modulo` without overflow.
56#[must_use]
57pub const fn mod_mul(a: u64, b: u64, modulo: u64) -> u64 {
58    ((a as u128 * b as u128) % modulo as u128) as u64
59}
60
61#[cfg(test)]
62mod test {
63    use super::*;
64
65    #[test]
66    fn check_isqrt() {
67        for x in 0..1_000_000 {
68            assert_eq!(isqrt(x), (x as f64).sqrt().floor() as u64);
69        }
70        assert_eq!(
71            f64::from(u32::MAX).sqrt().floor() as u64,
72            isqrt(u64::from(u32::MAX))
73        );
74        assert_eq!(isqrt(u64::MAX - 1), 4294967295);
75        assert_eq!(isqrt(u64::MAX), 4294967295);
76    }
77}