fixed_bigint/fixeduint/
isqrt_impl.rs

1// Copyright 2021 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Integer square root for FixedUInt.
16
17use super::{FixedUInt, MachineWord, const_set_bit};
18use crate::machineword::ConstMachineWord;
19use const_num_traits::Nct;
20use const_num_traits::{Isqrt, PrimBits, Zero};
21
22c0nst::c0nst! {
23    c0nst impl<T: [c0nst] ConstMachineWord + MachineWord, const N: usize> Isqrt for FixedUInt<T, N, Nct> {
24        type Output = Self;
25
26        fn isqrt(self) -> Self {
27            // Digit-by-digit (base-2) floor square root: largest r with
28            // r*r <= self, using only add/sub/shift — O(BIT_SIZE²) rather than
29            // the O(BIT_SIZE³) of a per-iteration `candidate * candidate` scan.
30            if <Self as Zero>::is_zero(&self) {
31                return <Self as Zero>::zero();
32            }
33
34            let mut num = self;
35            let mut res = <Self as Zero>::zero();
36
37            // Start at the largest power of four <= self.
38            let bit_len = Self::BIT_SIZE - PrimBits::leading_zeros(self) as usize;
39            let highest_even_bit = (bit_len - 1) & !1;
40            let mut bit = <Self as Zero>::zero();
41            const_set_bit(&mut bit.array, highest_even_bit);
42
43            while !<Self as Zero>::is_zero(&bit) {
44                let sum = res + bit;
45                if num >= sum {
46                    num -= sum;
47                    res = (res >> 1usize) + bit;
48                } else {
49                    res >>= 1usize;
50                }
51                bit >>= 2usize;
52            }
53
54            res
55        }
56    }
57
58    c0nst impl<T: [c0nst] ConstMachineWord + MachineWord, const N: usize> Isqrt for &FixedUInt<T, N, Nct> {
59        type Output = FixedUInt<T, N, Nct>;
60        fn isqrt(self) -> FixedUInt<T, N, Nct> {
61            <FixedUInt<T, N, Nct> as Isqrt>::isqrt(FixedUInt::from_array(self.array))
62        }
63    }
64}
65
66impl<T: ConstMachineWord + MachineWord, const N: usize> FixedUInt<T, N, Nct> {
67    /// Unsigned isqrt cannot fail; always returns `Some`.
68    pub fn checked_isqrt(self) -> Option<Self> {
69        Some(<Self as Isqrt>::isqrt(self))
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    #[cfg(feature = "num-traits")]
77    use num_traits::{CheckedAdd, CheckedMul};
78
79    #[test]
80    fn test_isqrt() {
81        type U16 = FixedUInt<u8, 2>;
82
83        // Perfect squares
84        assert_eq!(Isqrt::isqrt(U16::from(0u8)), U16::from(0u8));
85        assert_eq!(Isqrt::isqrt(U16::from(1u8)), U16::from(1u8));
86        assert_eq!(Isqrt::isqrt(U16::from(4u8)), U16::from(2u8));
87        assert_eq!(Isqrt::isqrt(U16::from(9u8)), U16::from(3u8));
88        assert_eq!(Isqrt::isqrt(U16::from(16u8)), U16::from(4u8));
89        assert_eq!(Isqrt::isqrt(U16::from(25u8)), U16::from(5u8));
90        assert_eq!(Isqrt::isqrt(U16::from(100u8)), U16::from(10u8));
91        assert_eq!(Isqrt::isqrt(U16::from(144u8)), U16::from(12u8));
92
93        // Non-perfect squares (floor)
94        assert_eq!(Isqrt::isqrt(U16::from(2u8)), U16::from(1u8));
95        assert_eq!(Isqrt::isqrt(U16::from(3u8)), U16::from(1u8));
96        assert_eq!(Isqrt::isqrt(U16::from(5u8)), U16::from(2u8));
97        assert_eq!(Isqrt::isqrt(U16::from(8u8)), U16::from(2u8));
98        assert_eq!(Isqrt::isqrt(U16::from(10u8)), U16::from(3u8));
99        assert_eq!(Isqrt::isqrt(U16::from(15u8)), U16::from(3u8));
100        assert_eq!(Isqrt::isqrt(U16::from(24u8)), U16::from(4u8));
101    }
102
103    #[test]
104    fn test_isqrt_larger_values() {
105        type U16 = FixedUInt<u8, 2>;
106
107        // Larger values
108        assert_eq!(Isqrt::isqrt(U16::from(10000u16)), U16::from(100u8));
109        assert_eq!(Isqrt::isqrt(U16::from(65535u16)), U16::from(255u8)); // sqrt(65535) = 255.998...
110        assert_eq!(Isqrt::isqrt(U16::from(65025u16)), U16::from(255u8)); // 255^2 = 65025
111    }
112
113    #[test]
114    fn test_checked_isqrt() {
115        type U16 = FixedUInt<u8, 2>;
116
117        // For unsigned, checked_isqrt always returns Some
118        assert_eq!(
119            FixedUInt::checked_isqrt(U16::from(0u8)),
120            Some(U16::from(0u8))
121        );
122        assert_eq!(
123            FixedUInt::checked_isqrt(U16::from(16u8)),
124            Some(U16::from(4u8))
125        );
126        assert_eq!(
127            FixedUInt::checked_isqrt(U16::from(17u8)),
128            Some(U16::from(4u8))
129        );
130    }
131
132    #[cfg(feature = "num-traits")]
133    #[test]
134    fn test_isqrt_correctness() {
135        type U16 = FixedUInt<u8, 2>;
136
137        // Verify r^2 <= n < (r+1)^2 for various values
138        for n in 0..=1000u16 {
139            let n_int = U16::from(n);
140            let r = Isqrt::isqrt(n_int);
141
142            // r^2 <= n
143            assert!(r * r <= n_int, "Failed: {:?}^2 > {}", r, n);
144
145            // (r+1)^2 > n - use checked arithmetic to handle potential overflow
146            if let Some(r_plus_1) = r.checked_add(&U16::from(1u8)) {
147                // If (r+1)^2 overflows, it's definitely > n since n fits in U16
148                if let Some(square) = r_plus_1.checked_mul(&r_plus_1) {
149                    assert!(square > n_int, "Failed: {:?}^2 <= {}", r_plus_1, n);
150                }
151            }
152            // If r+1 overflows, r is MAX, so (r+1)^2 > n also holds
153        }
154    }
155
156    #[cfg(feature = "num-traits")]
157    #[test]
158    fn test_isqrt_wider_types() {
159        // Test with wider word type to exercise cross-word bit-setting
160        type U32x2 = FixedUInt<u32, 2>;
161
162        // Perfect squares
163        assert_eq!(Isqrt::isqrt(U32x2::from(0u8)), U32x2::from(0u8));
164        assert_eq!(Isqrt::isqrt(U32x2::from(1u8)), U32x2::from(1u8));
165        assert_eq!(Isqrt::isqrt(U32x2::from(16u8)), U32x2::from(4u8));
166
167        // Larger values that span multiple bits
168        assert_eq!(Isqrt::isqrt(U32x2::from(1000000u32)), U32x2::from(1000u32));
169        assert_eq!(
170            Isqrt::isqrt(U32x2::from(0xFFFFFFFFu32)),
171            U32x2::from(0xFFFFu32)
172        );
173
174        // Test with u8x4 for different word boundary behavior
175        type U8x4 = FixedUInt<u8, 4>;
176        assert_eq!(Isqrt::isqrt(U8x4::from(65536u32)), U8x4::from(256u32));
177        assert_eq!(Isqrt::isqrt(U8x4::from(1000000u32)), U8x4::from(1000u32));
178
179        // Verify correctness for a range
180        for n in (0..=10000u32).step_by(100) {
181            let n_int = U32x2::from(n);
182            let r = Isqrt::isqrt(n_int);
183
184            // r^2 <= n
185            assert!(r * r <= n_int, "Failed: {:?}^2 > {} for U32x2", r, n);
186
187            // (r+1)^2 > n
188            if let Some(r_plus_1) = r.checked_add(&U32x2::from(1u8)) {
189                if let Some(square) = r_plus_1.checked_mul(&r_plus_1) {
190                    assert!(
191                        square > n_int,
192                        "Failed: {:?}^2 <= {} for U32x2",
193                        r_plus_1,
194                        n
195                    );
196                }
197            }
198        }
199    }
200
201    c0nst::c0nst! {
202        pub c0nst fn const_isqrt<T: [c0nst] ConstMachineWord + MachineWord, const N: usize>(
203            v: FixedUInt<T, N, Nct>,
204        ) -> FixedUInt<T, N, Nct> {
205            Isqrt::isqrt(v)
206        }
207        /// Const-callable parallel to `FixedUInt::checked_isqrt` (which
208        /// can't itself be `const fn` on an inherent impl, see the
209        /// shim's doc comment). External `CheckedIsqrt` is signed-only;
210        /// for unsigned types `isqrt` never fails, so this just lifts
211        /// the result into `Some`.
212        pub c0nst fn const_checked_isqrt<T: [c0nst] ConstMachineWord + MachineWord, const N: usize>(
213            v: FixedUInt<T, N, Nct>,
214        ) -> Option<FixedUInt<T, N, Nct>> {
215            Some(Isqrt::isqrt(v))
216        }
217    }
218
219    #[test]
220    fn test_const_isqrt() {
221        type U16 = FixedUInt<u8, 2>;
222
223        assert_eq!(const_isqrt(U16::from(16u8)), U16::from(4u8));
224        assert_eq!(const_isqrt(U16::from(100u8)), U16::from(10u8));
225        assert_eq!(const_checked_isqrt(U16::from(16u8)), Some(U16::from(4u8)));
226
227        #[cfg(feature = "nightly")]
228        {
229            const SIXTEEN: U16 = FixedUInt::from_array([16, 0]);
230            const RESULT: U16 = const_isqrt(SIXTEEN);
231            const CHECKED: Option<U16> = const_checked_isqrt(SIXTEEN);
232            assert_eq!(RESULT, FixedUInt::from_array([4, 0]));
233            assert!(CHECKED.is_some());
234        }
235    }
236}
fixed_bigint/fixeduint/isqrt_impl.rs

fixed_bigint/fixeduint/
isqrt_impl.rs
