bimm_contracts/
math.rs

1//! # Mathematical utilities.
2
3/// Find the exact integer nth root if it exists.
4///
5/// ## Arguments
6///
7/// - `value` - the value.
8/// - `exp` - the exponent.
9///
10/// ## Returns
11///
12/// Either the exact root (positive if possible) or None.
13pub fn maybe_iroot(value: isize, exp: usize) -> Option<isize> {
14    if exp == 1 {
15        return Some(value);
16    }
17    if exp == 0 {
18        return Some(1);
19    }
20    let (value, coeff) = if value >= 0 {
21        (value, 1)
22    } else if exp % 2 == 1 {
23        (-value, -1)
24    } else {
25        // No real root for negative value with even exponent
26        return None;
27    };
28    pos_maybe_iroot(value as usize, exp).map(|v| coeff * (v as isize))
29}
30
31/// Inner for `maybe_iroot`
32#[inline(always)]
33fn pos_maybe_iroot(value: usize, exp: usize) -> Option<usize> {
34    if exp == 1 {
35        return Some(value);
36    }
37    if exp == 0 {
38        return Some(1);
39    }
40    match value {
41        0 => Some(0),
42        1 => Some(1),
43        _ => {
44            let exp = exp as u32;
45
46            // Bisecting over the exclusive candidate range (lower, upper).
47            let mut lower = 1usize;
48
49            let mut upper = value.isqrt();
50            if exp == 2 && value == upper.pow(2) {
51                // Short-circuit for the builtin `isqrt` + squares.
52                return Some(upper);
53            }
54            upper += 1;
55
56            loop {
57                if lower + 1 >= upper {
58                    // No integer root found
59                    return None;
60                }
61                let candidate = (lower + upper) / 2;
62
63                match candidate.checked_pow(exp) {
64                    None => {
65                        // The candidate overflows; reduce the upper bound.
66                        upper = candidate - 1;
67                    }
68                    Some(pow) => {
69                        if pow == value {
70                            // Found the exact root
71                            return Some(candidate);
72                        } else if pow > value {
73                            // The candidate is too high, reduce the upper bound
74                            upper = candidate;
75                        } else {
76                            // The candidate is too low, increase the lower bound
77                            lower = candidate;
78                        }
79                    }
80                }
81            }
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn test_maybe_iroot() {
92        assert_eq!(maybe_iroot(0, 0), Some(1));
93        assert_eq!(maybe_iroot(0, 2), Some(0));
94
95        assert_eq!(maybe_iroot(4 * 4 * 4, 3), Some(4));
96
97        assert_eq!(maybe_iroot(1, 2), Some(1));
98        assert_eq!(maybe_iroot(4, 2), Some(2));
99        assert_eq!(maybe_iroot(8, 3), Some(2));
100        assert_eq!(maybe_iroot(27, 3), Some(3));
101
102        // exp == 1 should return the value itself
103        assert_eq!(maybe_iroot(5, 1), Some(5));
104        assert_eq!(maybe_iroot(-7, 1), Some(-7));
105
106        assert_eq!(maybe_iroot(-8, 3), Some(-2));
107        assert_eq!(maybe_iroot(-16, 4), None);
108        assert_eq!(maybe_iroot(16, 4), Some(2));
109        assert_eq!(maybe_iroot(-1, 3), Some(-1));
110        assert_eq!(maybe_iroot(-1, 2), None);
111
112        // Overflow case
113        let too_big = isize::MAX.isqrt() + 1;
114        assert_eq!(maybe_iroot(too_big, 2), None);
115        assert_eq!(maybe_iroot(too_big, 6), None);
116    }
117}
bimm_contracts/math.rs

bimm_contracts/
math.rs
