const_primes/
integer_math.rs

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