fixed_bigint/heapless/
isqrt.rs

1//! `const_num_traits::Isqrt` for `HeaplessBigInt<_, Nct>`.
2//!
3//! Digit-by-digit (base-2) floor square root, Nct-only (data-dependent
4//! comparisons). Uses only add/sub/shift — no per-iteration multiply — so it
5//! is O(width²) rather than the O(width³) of a `candidate * candidate` scan.
6//! `res`/`bit` are seeded at the operand width, and the `>> 1` / `>> 2` are
7//! sub-word shifts (never crossing a limb), so every value stays at `self.len`
8//! throughout and the result carries the operand width, bit-for-bit with the
9//! same-width `FixedUInt`.
10
11use super::HeaplessBigInt;
12use crate::MachineWord;
13use const_num_traits::{Isqrt, Nct, One, Zero};
14
15impl<T, const CAP: usize> Isqrt for HeaplessBigInt<T, CAP, Nct>
16where
17    T: MachineWord,
18{
19    type Output = Self;
20
21    fn isqrt(self) -> Self {
22        // Zero is its own root and already carries the operand width.
23        if <Self as Zero>::is_zero(&self) {
24            return self;
25        }
26
27        let width = self.len();
28        let one_w = <Self as One>::one().widened(width);
29        let mut num = self;
30        let mut res = Self::new_zero_with_len(width);
31
32        // Start at the largest power of four <= self: `1 << (highest even bit)`.
33        let highest_even_bit = (self.bit_length() - 1) & !1;
34        let mut bit = one_w << highest_even_bit;
35
36        while !<Self as Zero>::is_zero(&bit) {
37            let sum = res + bit;
38            if num >= sum {
39                num -= sum;
40                res = (res >> 1usize) + bit;
41            } else {
42                res >>= 1usize;
43            }
44            bit >>= 2usize;
45        }
46        res
47    }
48}
49
50// `&Self` mirror so `(&h).isqrt()` resolves without an explicit copy.
51impl<T, const CAP: usize> Isqrt for &HeaplessBigInt<T, CAP, Nct>
52where
53    T: MachineWord,
54{
55    type Output = HeaplessBigInt<T, CAP, Nct>;
56    fn isqrt(self) -> Self::Output {
57        <HeaplessBigInt<T, CAP, Nct> as Isqrt>::isqrt(*self)
58    }
59}
60
61impl<T, const CAP: usize> HeaplessBigInt<T, CAP, Nct>
62where
63    T: MachineWord,
64{
65    /// Unsigned isqrt cannot fail; always `Some`. Parallels
66    /// `FixedUInt::checked_isqrt` and the signed-only external
67    /// `CheckedIsqrt`.
68    #[must_use]
69    pub fn checked_isqrt(self) -> Option<Self> {
70        Some(<Self as Isqrt>::isqrt(self))
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::HeaplessBigInt;
77    use const_num_traits::Isqrt;
78
79    type H = HeaplessBigInt<u8, 8>;
80
81    #[test]
82    fn isqrt_values() {
83        for (n, r) in [
84            (0u16, 0u16),
85            (1, 1),
86            (4, 2),
87            (15, 3),
88            (16, 4),
89            (24, 4),
90            (10000, 100),
91            (65535, 255),
92        ] {
93            assert_eq!(Isqrt::isqrt(H::from(n)), H::from(r), "isqrt({n})");
94        }
95    }
96
97    // Matches a reference integer isqrt across a range, exercising the
98    // digit-by-digit steps.
99    #[test]
100    fn isqrt_correctness_range() {
101        fn ref_isqrt(n: u32) -> u32 {
102            let mut r = 0u32;
103            while (r + 1) * (r + 1) <= n {
104                r += 1;
105            }
106            r
107        }
108        for n in 0u16..=2000 {
109            let expected = ref_isqrt(u32::from(n)) as u16;
110            assert_eq!(Isqrt::isqrt(H::from(n)), H::from(expected), "isqrt({n})");
111        }
112    }
113
114    // The result carries the operand width, not the minimal magnitude width.
115    #[test]
116    fn isqrt_preserves_width() {
117        let n = H::from(10000u16).widened(8);
118        let r = Isqrt::isqrt(n);
119        assert_eq!(r, H::from(100u8));
120        assert_eq!(r.len(), 8);
121
122        // Zero keeps its width too.
123        let z = H::new_zero_with_len(8);
124        assert_eq!(Isqrt::isqrt(z).len(), 8);
125    }
126
127    #[test]
128    fn byref_matches_value() {
129        let a = H::from(10000u16);
130        assert_eq!(Isqrt::isqrt(&a), Isqrt::isqrt(a));
131    }
132}
fixed_bigint/heapless/isqrt.rs

fixed_bigint/heapless/
isqrt.rs
