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            // Bit-by-bit algorithm for integer square root.
28            // Returns the largest r such that r * r <= self.
29            if <Self as Zero>::is_zero(&self) {
30                return <Self as Zero>::zero();
31            }
32
33            let mut result = <Self as Zero>::zero();
34
35            let bit_len = Self::BIT_SIZE - PrimBits::leading_zeros(self) as usize;
36            let start_bit = bit_len.div_ceil(2);
37
38            let mut bit_pos = start_bit;
39            while bit_pos > 0 {
40                bit_pos -= 1;
41
42                let mut candidate = result;
43                const_set_bit(&mut candidate.array, bit_pos);
44
45                let square = candidate * candidate;
46                if square <= self {
47                    result = candidate;
48                }
49            }
50
51            result
52        }
53    }
54
55    c0nst impl<T: [c0nst] ConstMachineWord + MachineWord, const N: usize> Isqrt for &FixedUInt<T, N, Nct> {
56        type Output = FixedUInt<T, N, Nct>;
57        fn isqrt(self) -> FixedUInt<T, N, Nct> {
58            <FixedUInt<T, N, Nct> as Isqrt>::isqrt(*self)
59        }
60    }
61}
62
63impl<T: ConstMachineWord + MachineWord, const N: usize> FixedUInt<T, N, Nct> {
64    /// Unsigned isqrt cannot fail; always returns `Some`.
65    pub fn checked_isqrt(self) -> Option<Self> {
66        Some(<Self as Isqrt>::isqrt(self))
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73    #[cfg(feature = "num-traits")]
74    use num_traits::{CheckedAdd, CheckedMul};
75
76    #[test]
77    fn test_isqrt() {
78        type U16 = FixedUInt<u8, 2>;
79
80        // Perfect squares
81        assert_eq!(Isqrt::isqrt(U16::from(0u8)), U16::from(0u8));
82        assert_eq!(Isqrt::isqrt(U16::from(1u8)), U16::from(1u8));
83        assert_eq!(Isqrt::isqrt(U16::from(4u8)), U16::from(2u8));
84        assert_eq!(Isqrt::isqrt(U16::from(9u8)), U16::from(3u8));
85        assert_eq!(Isqrt::isqrt(U16::from(16u8)), U16::from(4u8));
86        assert_eq!(Isqrt::isqrt(U16::from(25u8)), U16::from(5u8));
87        assert_eq!(Isqrt::isqrt(U16::from(100u8)), U16::from(10u8));
88        assert_eq!(Isqrt::isqrt(U16::from(144u8)), U16::from(12u8));
89
90        // Non-perfect squares (floor)
91        assert_eq!(Isqrt::isqrt(U16::from(2u8)), U16::from(1u8));
92        assert_eq!(Isqrt::isqrt(U16::from(3u8)), U16::from(1u8));
93        assert_eq!(Isqrt::isqrt(U16::from(5u8)), U16::from(2u8));
94        assert_eq!(Isqrt::isqrt(U16::from(8u8)), U16::from(2u8));
95        assert_eq!(Isqrt::isqrt(U16::from(10u8)), U16::from(3u8));
96        assert_eq!(Isqrt::isqrt(U16::from(15u8)), U16::from(3u8));
97        assert_eq!(Isqrt::isqrt(U16::from(24u8)), U16::from(4u8));
98    }
99
100    #[test]
101    fn test_isqrt_larger_values() {
102        type U16 = FixedUInt<u8, 2>;
103
104        // Larger values
105        assert_eq!(Isqrt::isqrt(U16::from(10000u16)), U16::from(100u8));
106        assert_eq!(Isqrt::isqrt(U16::from(65535u16)), U16::from(255u8)); // sqrt(65535) = 255.998...
107        assert_eq!(Isqrt::isqrt(U16::from(65025u16)), U16::from(255u8)); // 255^2 = 65025
108    }
109
110    #[test]
111    fn test_checked_isqrt() {
112        type U16 = FixedUInt<u8, 2>;
113
114        // For unsigned, checked_isqrt always returns Some
115        assert_eq!(
116            FixedUInt::checked_isqrt(U16::from(0u8)),
117            Some(U16::from(0u8))
118        );
119        assert_eq!(
120            FixedUInt::checked_isqrt(U16::from(16u8)),
121            Some(U16::from(4u8))
122        );
123        assert_eq!(
124            FixedUInt::checked_isqrt(U16::from(17u8)),
125            Some(U16::from(4u8))
126        );
127    }
128
129    #[cfg(feature = "num-traits")]
130    #[test]
131    fn test_isqrt_correctness() {
132        type U16 = FixedUInt<u8, 2>;
133
134        // Verify r^2 <= n < (r+1)^2 for various values
135        for n in 0..=1000u16 {
136            let n_int = U16::from(n);
137            let r = Isqrt::isqrt(n_int);
138
139            // r^2 <= n
140            assert!(r * r <= n_int, "Failed: {:?}^2 > {}", r, n);
141
142            // (r+1)^2 > n - use checked arithmetic to handle potential overflow
143            if let Some(r_plus_1) = r.checked_add(&U16::from(1u8)) {
144                // If (r+1)^2 overflows, it's definitely > n since n fits in U16
145                if let Some(square) = r_plus_1.checked_mul(&r_plus_1) {
146                    assert!(square > n_int, "Failed: {:?}^2 <= {}", r_plus_1, n);
147                }
148            }
149            // If r+1 overflows, r is MAX, so (r+1)^2 > n also holds
150        }
151    }
152
153    #[cfg(feature = "num-traits")]
154    #[test]
155    fn test_isqrt_wider_types() {
156        // Test with wider word type to exercise cross-word bit-setting
157        type U32x2 = FixedUInt<u32, 2>;
158
159        // Perfect squares
160        assert_eq!(Isqrt::isqrt(U32x2::from(0u8)), U32x2::from(0u8));
161        assert_eq!(Isqrt::isqrt(U32x2::from(1u8)), U32x2::from(1u8));
162        assert_eq!(Isqrt::isqrt(U32x2::from(16u8)), U32x2::from(4u8));
163
164        // Larger values that span multiple bits
165        assert_eq!(Isqrt::isqrt(U32x2::from(1000000u32)), U32x2::from(1000u32));
166        assert_eq!(
167            Isqrt::isqrt(U32x2::from(0xFFFFFFFFu32)),
168            U32x2::from(0xFFFFu32)
169        );
170
171        // Test with u8x4 for different word boundary behavior
172        type U8x4 = FixedUInt<u8, 4>;
173        assert_eq!(Isqrt::isqrt(U8x4::from(65536u32)), U8x4::from(256u32));
174        assert_eq!(Isqrt::isqrt(U8x4::from(1000000u32)), U8x4::from(1000u32));
175
176        // Verify correctness for a range
177        for n in (0..=10000u32).step_by(100) {
178            let n_int = U32x2::from(n);
179            let r = Isqrt::isqrt(n_int);
180
181            // r^2 <= n
182            assert!(r * r <= n_int, "Failed: {:?}^2 > {} for U32x2", r, n);
183
184            // (r+1)^2 > n
185            if let Some(r_plus_1) = r.checked_add(&U32x2::from(1u8)) {
186                if let Some(square) = r_plus_1.checked_mul(&r_plus_1) {
187                    assert!(
188                        square > n_int,
189                        "Failed: {:?}^2 <= {} for U32x2",
190                        r_plus_1,
191                        n
192                    );
193                }
194            }
195        }
196    }
197
198    c0nst::c0nst! {
199        pub c0nst fn const_isqrt<T: [c0nst] ConstMachineWord + MachineWord, const N: usize>(
200            v: FixedUInt<T, N, Nct>,
201        ) -> FixedUInt<T, N, Nct> {
202            Isqrt::isqrt(v)
203        }
204        /// Const-callable parallel to `FixedUInt::checked_isqrt` (which
205        /// can't itself be `const fn` on an inherent impl, see the
206        /// shim's doc comment). External `CheckedIsqrt` is signed-only;
207        /// for unsigned types `isqrt` never fails, so this just lifts
208        /// the result into `Some`.
209        pub c0nst fn const_checked_isqrt<T: [c0nst] ConstMachineWord + MachineWord, const N: usize>(
210            v: FixedUInt<T, N, Nct>,
211        ) -> Option<FixedUInt<T, N, Nct>> {
212            Some(Isqrt::isqrt(v))
213        }
214    }
215
216    #[test]
217    fn test_const_isqrt() {
218        type U16 = FixedUInt<u8, 2>;
219
220        assert_eq!(const_isqrt(U16::from(16u8)), U16::from(4u8));
221        assert_eq!(const_isqrt(U16::from(100u8)), U16::from(10u8));
222        assert_eq!(const_checked_isqrt(U16::from(16u8)), Some(U16::from(4u8)));
223
224        #[cfg(feature = "nightly")]
225        {
226            const SIXTEEN: U16 = FixedUInt::from_array([16, 0]);
227            const RESULT: U16 = const_isqrt(SIXTEEN);
228            const CHECKED: Option<U16> = const_checked_isqrt(SIXTEEN);
229            assert_eq!(RESULT, FixedUInt::from_array([4, 0]));
230            assert!(CHECKED.is_some());
231        }
232    }
233}
fixed_bigint/fixeduint/isqrt_impl.rs

fixed_bigint/fixeduint/
isqrt_impl.rs
