Skip to main content

dashu_base/math/
root.rs

1use super::{CubicRoot, SquareRoot};
2use crate::{CubicRootRem, NormalizedRootRem, SquareRootRem};
3
4impl SquareRoot for u8 {
5    type Output = u8;
6
7    #[inline]
8    fn sqrt(&self) -> Self::Output {
9        self.sqrt_rem().0
10    }
11}
12
13impl CubicRoot for u8 {
14    type Output = u8;
15
16    #[inline]
17    fn cbrt(&self) -> Self::Output {
18        self.cbrt_rem().0
19    }
20}
21
22// f64::sqrt is reliable for floor(sqrt(x)) when x < (2^26+1)^2.
23// Beyond that, IEEE 754 rounding can push the result across an integer
24// boundary (e.g. f64::sqrt(67108865^2 - 1) rounds up to 67108865.0).
25#[cfg(feature = "std")]
26const F64_SQRT_THRESHOLD: u128 = 4503599761588224; // (2^26 + 1)^2 - 1
27
28// When std is available, use native float sqrt for each primitive type.
29// For u16, f32 is sufficient (24-bit mantissa covers all u16 values exactly).
30// For u32/u64/u128, f64 is used with the [F64_SQRT_THRESHOLD] guard.
31
32#[cfg(feature = "std")]
33impl SquareRoot for u16 {
34    type Output = u8;
35
36    #[inline]
37    fn sqrt(&self) -> u8 {
38        if *self == 0 {
39            return 0;
40        }
41        (*self as f32).sqrt().floor() as u8
42    }
43}
44
45#[cfg(feature = "std")]
46impl SquareRoot for u32 {
47    type Output = u16;
48
49    #[inline]
50    fn sqrt(&self) -> u16 {
51        if *self == 0 {
52            return 0;
53        }
54        (*self as f64).sqrt().floor() as u16
55    }
56}
57
58#[cfg(feature = "std")]
59impl SquareRoot for u64 {
60    type Output = u32;
61
62    #[inline]
63    fn sqrt(&self) -> u32 {
64        if *self == 0 {
65            return 0;
66        }
67        // u64 values above the threshold fall back to the bitwise algorithm.
68        if (*self as u128) < F64_SQRT_THRESHOLD {
69            return (*self as f64).sqrt().floor() as u32;
70        }
71
72        let shift = self.leading_zeros() & !1; // make sure shift is divisible by 2
73        let (root, _) = (self << shift).normalized_sqrt_rem();
74        root >> (shift / 2)
75    }
76}
77
78#[cfg(feature = "std")]
79impl SquareRoot for u128 {
80    type Output = u64;
81
82    #[inline]
83    fn sqrt(&self) -> u64 {
84        if *self == 0 {
85            return 0;
86        }
87        if *self < F64_SQRT_THRESHOLD {
88            return (*self as f64).sqrt().floor() as u64;
89        }
90
91        let shift = self.leading_zeros() & !1; // make sure shift is divisible by 2
92        let (root, _) = (self << shift).normalized_sqrt_rem();
93        root >> (shift / 2)
94    }
95}
96
97// When std is not available, implement SquareRoot through a macro using the
98// bitwise normalization approach.
99#[cfg(not(feature = "std"))]
100macro_rules! impl_sqrt_using_rootrem {
101    ($t:ty, $half:ty) => {
102        impl SquareRoot for $t {
103            type Output = $half;
104
105            #[inline]
106            fn sqrt(&self) -> $half {
107                if *self == 0 {
108                    return 0;
109                }
110
111                // normalize the input and call the normalized subroutine
112                let shift = self.leading_zeros() & !1; // make sure shift is divisible by 2
113                let (root, _) = (self << shift).normalized_sqrt_rem();
114                root >> (shift / 2)
115            }
116        }
117    };
118}
119
120#[cfg(not(feature = "std"))]
121impl_sqrt_using_rootrem!(u16, u8);
122#[cfg(not(feature = "std"))]
123impl_sqrt_using_rootrem!(u32, u16);
124#[cfg(not(feature = "std"))]
125impl_sqrt_using_rootrem!(u64, u32);
126#[cfg(not(feature = "std"))]
127impl_sqrt_using_rootrem!(u128, u64);
128
129// CubicRoot for all types above u8 always uses the bitwise approach.
130macro_rules! impl_cbrt_using_cbrtrem {
131    ($t:ty, $half:ty) => {
132        impl CubicRoot for $t {
133            type Output = $half;
134
135            #[inline]
136            fn cbrt(&self) -> $half {
137                if *self == 0 {
138                    return 0;
139                }
140
141                // normalize the input and call the normalized subroutine
142                let mut shift = self.leading_zeros();
143                shift -= shift % 3; // make sure shift is divisible by 3
144                let (root, _) = (self << shift).normalized_cbrt_rem();
145                root >> (shift / 3)
146            }
147        }
148    };
149}
150
151impl_cbrt_using_cbrtrem!(u16, u8);
152impl_cbrt_using_cbrtrem!(u32, u16);
153impl_cbrt_using_cbrtrem!(u64, u32);
154impl_cbrt_using_cbrtrem!(u128, u64);