ruint/
root.rs

Help
1#![cfg(feature = "std")]
2
3use crate::Uint;
4use core::cmp::{min, Ordering};
5
6impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
7    /// Computes the floor of the `degree`-th root of the number.
8    ///
9    /// $$
10    /// \floor{\sqrt[\mathtt{degree}]{\mathtt{self}}}
11    /// $$
12    ///
13    /// # Panics
14    ///
15    /// Panics if `degree` is zero.
16    ///
17    /// # Examples
18    ///
19    /// ```
20    /// # use ruint::{Uint, uint, aliases::*};
21    /// # uint!{
22    /// assert_eq!(0_U64.root(2), 0_U64);
23    /// assert_eq!(1_U64.root(63), 1_U64);
24    /// assert_eq!(0x0032da8b0f88575d_U63.root(64), 1_U63);
25    /// assert_eq!(0x1756800000000000_U63.root(34), 3_U63);
26    /// # }
27    /// ```
28    #[inline]
29    #[must_use]
30    pub fn root(self, degree: usize) -> Self {
31        assert!(degree > 0, "degree must be greater than zero");
32
33        // Handle zero case (including BITS == 0).
34        if self.is_zero() {
35            return Self::ZERO;
36        }
37
38        // Handle case where `degree > Self::BITS`.
39        if degree >= Self::BITS {
40            return Self::from(1);
41        }
42
43        // Handle case where `degree == 1`.
44        if degree == 1 {
45            return self;
46        }
47
48        // Create a first guess.
49        // Root should be less than the value, so approx_pow2 should always succeed.
50        #[allow(clippy::cast_precision_loss)] // Approximation is good enough.
51        #[allow(clippy::cast_sign_loss)] // Result should be positive.
52        let mut result = Self::approx_pow2(self.approx_log2() / degree as f64).unwrap();
53
54        let deg_m1 = Self::from(degree - 1);
55
56        // Iterate using Newton's method
57        // See <https://en.wikipedia.org/wiki/Integer_square_root#Algorithm_using_Newton's_method>
58        // See <https://gmplib.org/manual/Nth-Root-Algorithm>
59        let mut decreasing = false;
60        loop {
61            // OPT: This could benefit from single-limb multiplication
62            // and division.
63            //
64            // OPT: The division can be turned into bit-shifts when the degree is a power of
65            // two.
66            let division = result
67                .checked_pow(deg_m1)
68                .map_or(Self::ZERO, |power| self / power);
69            let iter = (division + deg_m1 * result) / Self::from(degree);
70            match (decreasing, iter.cmp(&result)) {
71                // Stop when we hit fix point or stop decreasing.
72                (_, Ordering::Equal) | (true, Ordering::Greater) => break result,
73
74                // When `degree` is high and the initial guess is less than or equal to the
75                // (small) true result, it takes a long time to converge. Example:
76                // 0x215f07147d573ef203e1f268ab1516d3f294619db820c5dfd0b334e4d06320b7_U256.
77                // root(196) takes 5918 iterations to converge from the initial guess of `2`.
78                // to the final result of `2`. This is because after the first iteration
79                // it jumps to `1533576856264507`. To fix this we cap the increase at `2x`.
80                // Once `result` exceeds the true result, it will converge downwards.
81                (false, Ordering::Greater) => result = min(iter, result.saturating_shl(1)),
82
83                // Converging downwards.
84                (_, Ordering::Less) => {
85                    decreasing = true;
86                    result = iter;
87                }
88            }
89        }
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use crate::{const_for, nlimbs};
97    use proptest::proptest;
98
99    #[test]
100    #[allow(clippy::absurd_extreme_comparisons)] // From macro.
101    fn test_root() {
102        const_for!(BITS in SIZES if (BITS > 3) {
103            const LIMBS: usize = nlimbs(BITS);
104            type U = Uint<BITS, LIMBS>;
105            proptest!(|(value: U, degree in 1_usize..=5)| {
106                let root = value.root(degree);
107                let lower = root.pow(U::from(degree));
108                assert!(value >= lower);
109                let upper = root
110                    .checked_add(U::from(1))
111                    .and_then(|n| n.checked_pow(U::from(degree)));
112                if let Some(upper) = upper {
113                   assert!(value < upper);
114                }
115            });
116        });
117    }
118
119    #[test]
120    #[allow(clippy::absurd_extreme_comparisons)] // From macro.
121    #[allow(clippy::reversed_empty_ranges)] // From macro.
122    fn test_root_large() {
123        const_for!(BITS in SIZES if (BITS > 3) {
124            const LIMBS: usize = nlimbs(BITS);
125            type U = Uint<BITS, LIMBS>;
126            proptest!(|(value: U, degree in 1_usize..=BITS)| {
127                let root = value.root(degree);
128                let lower = root.pow(U::from(degree));
129                assert!(value >= lower);
130                let upper = root
131                    .checked_add(U::from(1))
132                    .and_then(|n| n.checked_pow(U::from(degree)));
133                if let Some(upper) = upper {
134                   assert!(value < upper);
135                }
136            });
137        });
138    }
139}
ruint/root.rs

ruint/
root.rs
