fixed_bigint/fixeduint/
roots_impl.rs

1use crate::fixeduint::FixedUInt;
2use crate::machineword::MachineWord;
3use const_num_traits::{CheckedPow, Nct};
4use num_integer::Roots;
5use num_traits::{FromPrimitive, One, Zero};
6
7impl<T: MachineWord, const N: usize> Roots for FixedUInt<T, N, Nct> {
8    fn nth_root(&self, n: u32) -> Self {
9        if n == 0 {
10            panic!("nth_root: n must be non-zero");
11        }
12
13        if self.is_zero() {
14            return Self::zero();
15        }
16
17        if self.is_one() || n == 1 {
18            return *self;
19        }
20
21        let bit_len = self.bit_length();
22        if n > bit_len {
23            return Self::one();
24        }
25
26        // Initial guess: use ceiling(bit_len / n) for overestimate
27        let initial_exp = bit_len.div_ceil(n).max(1);
28        let mut x = Self::one() << (initial_exp as usize);
29
30        // Constants using FromPrimitive
31        let n_val = Self::from_u32(n).expect("n too large for FixedUInt");
32        let n_minus_1 = Self::from_u32(n - 1).expect("n too large for FixedUInt");
33
34        // Newton's method iteration. Power evaluations go through `checked_pow`:
35        // an over-width `x^k` would panic (Nct multiply overflow), so overflow
36        // is treated as "arbitrarily large" — the quotient `self / x^(n-1)` is 0
37        // and the upper probe `(x+1)^n` counts as greater than `self`.
38        loop {
39            let quotient = match CheckedPow::checked_pow(x, n - 1) {
40                Some(p) if p.is_zero() => break,
41                Some(p) => *self / p,
42                None => Self::zero(),
43            };
44
45            let numerator = x * n_minus_1 + quotient;
46            let x_new = numerator / n_val;
47
48            if x_new >= x {
49                break;
50            }
51
52            x = x_new;
53        }
54
55        // Final adjustment to ensure r^n <= self < (r+1)^n. `self` fits the
56        // width, so an over-width `x^n` (None) is necessarily greater.
57        while CheckedPow::checked_pow(x, n).is_none_or(|p| p > *self) {
58            x -= Self::one();
59        }
60
61        let mut x_plus_one = x + Self::one();
62        while CheckedPow::checked_pow(x_plus_one, n).is_some_and(|p| p <= *self) {
63            x += Self::one();
64            x_plus_one = x + Self::one();
65        }
66
67        x
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use num_integer::Roots;
75    use num_traits::One;
76
77    #[test]
78    fn test_sqrt_basic() {
79        type TestInt = FixedUInt<u32, 2>;
80
81        assert_eq!(TestInt::from(0u8).sqrt(), TestInt::from(0u8));
82        assert_eq!(TestInt::from(1u8).sqrt(), TestInt::from(1u8));
83        assert_eq!(TestInt::from(4u8).sqrt(), TestInt::from(2u8));
84        assert_eq!(TestInt::from(9u8).sqrt(), TestInt::from(3u8));
85        assert_eq!(TestInt::from(16u8).sqrt(), TestInt::from(4u8));
86        assert_eq!(TestInt::from(25u8).sqrt(), TestInt::from(5u8));
87
88        // Test non-perfect squares
89        assert_eq!(TestInt::from(2u8).sqrt(), TestInt::from(1u8));
90        assert_eq!(TestInt::from(3u8).sqrt(), TestInt::from(1u8));
91        assert_eq!(TestInt::from(8u8).sqrt(), TestInt::from(2u8));
92        assert_eq!(TestInt::from(15u8).sqrt(), TestInt::from(3u8));
93        assert_eq!(TestInt::from(24u8).sqrt(), TestInt::from(4u8));
94    }
95
96    #[test]
97    fn test_cbrt_basic() {
98        type TestInt = FixedUInt<u32, 2>;
99
100        assert_eq!(TestInt::from(0u8).cbrt(), TestInt::from(0u8));
101        assert_eq!(TestInt::from(1u8).cbrt(), TestInt::from(1u8));
102        assert_eq!(TestInt::from(8u8).cbrt(), TestInt::from(2u8));
103        assert_eq!(TestInt::from(27u8).cbrt(), TestInt::from(3u8));
104        assert_eq!(TestInt::from(64u8).cbrt(), TestInt::from(4u8));
105        assert_eq!(TestInt::from(125u8).cbrt(), TestInt::from(5u8));
106
107        // Test non-perfect cubes
108        assert_eq!(TestInt::from(2u8).cbrt(), TestInt::from(1u8));
109        assert_eq!(TestInt::from(7u8).cbrt(), TestInt::from(1u8));
110        assert_eq!(TestInt::from(26u8).cbrt(), TestInt::from(2u8));
111        assert_eq!(TestInt::from(63u8).cbrt(), TestInt::from(3u8));
112    }
113
114    #[test]
115    fn test_nth_root_basic() {
116        type TestInt = FixedUInt<u32, 2>;
117
118        // Test 4th roots
119        assert_eq!(TestInt::from(16u8).nth_root(4), TestInt::from(2u8));
120        assert_eq!(TestInt::from(81u8).nth_root(4), TestInt::from(3u8));
121        assert_eq!(TestInt::from(15u8).nth_root(4), TestInt::from(1u8));
122        assert_eq!(TestInt::from(80u8).nth_root(4), TestInt::from(2u8));
123
124        // Test 5th roots
125        assert_eq!(TestInt::from(32u8).nth_root(5), TestInt::from(2u8));
126        assert_eq!(TestInt::from(243u8).nth_root(5), TestInt::from(3u8));
127        assert_eq!(TestInt::from(31u8).nth_root(5), TestInt::from(1u8));
128
129        // Test n=1 (should return self)
130        assert_eq!(TestInt::from(42u8).nth_root(1), TestInt::from(42u8));
131        assert_eq!(TestInt::from(123u8).nth_root(1), TestInt::from(123u8));
132    }
133
134    #[test]
135    fn test_nth_root_edge_cases() {
136        type TestInt = FixedUInt<u32, 2>;
137
138        // Test with 0 and 1
139        assert_eq!(TestInt::from(0u8).nth_root(2), TestInt::from(0u8));
140        assert_eq!(TestInt::from(1u8).nth_root(2), TestInt::from(1u8));
141        assert_eq!(TestInt::from(0u8).nth_root(10), TestInt::from(0u8));
142        assert_eq!(TestInt::from(1u8).nth_root(10), TestInt::from(1u8));
143
144        // Test with large n (should return 1 for numbers > 1)
145        assert_eq!(TestInt::from(2u8).nth_root(100), TestInt::from(1u8));
146        assert_eq!(TestInt::from(1000u16).nth_root(50), TestInt::from(1u8));
147    }
148
149    #[test]
150    #[should_panic(expected = "nth_root: n must be non-zero")]
151    fn test_nth_root_zero_n() {
152        let x = FixedUInt::<u32, 2>::from(16u8);
153        x.nth_root(0);
154    }
155
156    #[test]
157    fn test_root_correctness() {
158        type TestInt = FixedUInt<u32, 2>;
159
160        // Test that r^n <= x < (r+1)^n for various cases
161        for x in 1..=100u16 {
162            let x_int = TestInt::from(x);
163
164            // Test square root
165            let sqrt_x = x_int.sqrt();
166            assert!(sqrt_x.pow(2) <= x_int);
167            assert!((sqrt_x + TestInt::one()).pow(2) > x_int);
168
169            // Test cube root
170            let cbrt_x = x_int.cbrt();
171            assert!(cbrt_x.pow(3) <= x_int);
172            assert!((cbrt_x + TestInt::one()).pow(3) > x_int);
173
174            // Test 4th root
175            let root4_x = x_int.nth_root(4);
176            assert!(root4_x.pow(4) <= x_int);
177            assert!((root4_x + TestInt::one()).pow(4) > x_int);
178        }
179    }
180
181    // sqrt of the max value when the type width == value width: the upper
182    // probe (x+1)^2 = 2^BITS overflows and used to panic in the Nct multiply.
183    // checked_pow now reads that overflow as "> self".
184    #[test]
185    fn sqrt_of_max_does_not_overflow_panic() {
186        assert_eq!(
187            FixedUInt::<u32, 1>::from(u32::MAX).sqrt(),
188            FixedUInt::<u32, 1>::from(0xFFFFu32)
189        );
190        assert_eq!(
191            FixedUInt::<u8, 1>::from(255u8).sqrt(),
192            FixedUInt::<u8, 1>::from(15u8)
193        );
194        // A perfect square at the top of the range, and its neighbours.
195        assert_eq!(
196            FixedUInt::<u16, 1>::from(65025u16).sqrt(),
197            FixedUInt::<u16, 1>::from(255u16)
198        ); // 255^2
199        assert_eq!(
200            FixedUInt::<u16, 1>::from(65535u16).sqrt(),
201            FixedUInt::<u16, 1>::from(255u16)
202        );
203    }
204}
fixed_bigint/fixeduint/roots_impl.rs

fixed_bigint/fixeduint/
roots_impl.rs
