gpl_math/
approximations.rs

1//! Approximation calculations
2
3use {
4    num_traits::{CheckedShl, CheckedShr, PrimInt},
5    std::cmp::Ordering,
6};
7
8/// Calculate square root of the given number
9///
10/// Code lovingly adapted from the excellent work at:
11///
12/// <https://github.com/derekdreery/integer-sqrt-rs>
13///
14/// The algorithm is based on the implementation in:
15///
16/// <https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)>
17pub fn sqrt<T: PrimInt + CheckedShl + CheckedShr>(radicand: T) -> Option<T> {
18    match radicand.cmp(&T::zero()) {
19        Ordering::Less => return None,             // fail for less than 0
20        Ordering::Equal => return Some(T::zero()), // do nothing for 0
21        _ => {}
22    }
23
24    // Compute bit, the largest power of 4 <= n
25    let max_shift: u32 = T::zero().leading_zeros() - 1;
26    let shift: u32 = (max_shift - radicand.leading_zeros()) & !1;
27    let mut bit = T::one().checked_shl(shift)?;
28
29    let mut n = radicand;
30    let mut result = T::zero();
31    while bit != T::zero() {
32        let result_with_bit = result.checked_add(&bit)?;
33        if n >= result_with_bit {
34            n = n.checked_sub(&result_with_bit)?;
35            result = result.checked_shr(1)?.checked_add(&bit)?;
36        } else {
37            result = result.checked_shr(1)?;
38        }
39        bit = bit.checked_shr(2)?;
40    }
41    Some(result)
42}
43
44#[cfg(test)]
45mod tests {
46    use {super::*, proptest::prelude::*};
47
48    fn check_square_root(radicand: u128) {
49        let root = sqrt(radicand).unwrap();
50        let lower_bound = root.saturating_sub(1).checked_pow(2).unwrap();
51        let upper_bound = root.checked_add(1).unwrap().checked_pow(2).unwrap();
52        assert!(radicand as u128 <= upper_bound);
53        assert!(radicand as u128 >= lower_bound);
54    }
55
56    #[test]
57    fn test_square_root_min_max() {
58        let test_roots = [0, u64::MAX];
59        for i in test_roots.iter() {
60            check_square_root(*i as u128);
61        }
62    }
63
64    proptest! {
65        #[test]
66        fn test_square_root(a in 0..u64::MAX) {
67            check_square_root(a as u128);
68        }
69    }
70}
gpl_math/approximations.rs

gpl_math/
approximations.rs
