malachite_base/num/factorization/
is_square.rs

1// Copyright © 2025 William Youmans
2//
3// Uses code adopted from the FLINT Library.
4//
5//      Copyright © 2009 William Hart
6//
7// This file is part of Malachite.
8//
9// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
10// Lesser General Public License (LGPL as published by the Free Software Foundation; either version
11// 3 of the License, or (at your option any later version. See <https://www.gnu.org/licenses/>.
12
13use crate::num::arithmetic::traits::{CheckedSqrt, FloorSqrt, Square};
14use crate::num::basic::integers::PrimitiveInt;
15use crate::num::conversion::traits::{SplitInHalf, WrappingFrom};
16use crate::num::factorization::traits::IsSquare;
17
18const IS_SQUARE_MOD64: [bool; 64] = [
19    true, true, false, false, true, false, false, false, false, true, false, false, false, false,
20    false, false, true, true, false, false, false, false, false, false, false, true, false, false,
21    false, false, false, false, false, true, false, false, true, false, false, false, false, true,
22    false, false, false, false, false, false, false, true, false, false, false, false, false,
23    false, false, true, false, false, false, false, false, false,
24];
25
26const IS_SQUARE_MOD65: [bool; 65] = [
27    true, true, false, false, true, false, false, false, false, true, true, false, false, false,
28    true, false, true, false, false, false, false, false, false, false, false, true, true, false,
29    false, true, true, false, false, false, false, true, true, false, false, true, true, false,
30    false, false, false, false, false, false, false, true, false, true, false, false, false, true,
31    true, false, false, false, false, true, false, false, true,
32];
33
34const IS_SQUARE_MOD63: [bool; 63] = [
35    true, true, false, false, true, false, false, true, false, true, false, false, false, false,
36    true, false, true, false, true, false, false, false, true, false, false, true, false, false,
37    true, false, false, false, false, false, false, true, true, true, false, false, false, false,
38    false, true, false, false, true, false, false, true, false, false, false, false, false, false,
39    true, false, true, false, false, false, false,
40];
41
42// This is n_is_square when FLINT64 is false, from ulong_extras/is_square.c, FLINT 3.1.2.
43fn is_square_u64(x: u64) -> bool {
44    IS_SQUARE_MOD64[(x % 64) as usize]
45        && IS_SQUARE_MOD63[(x % 63) as usize]
46        && IS_SQUARE_MOD65[(x % 65) as usize]
47        && x.floor_sqrt().square() == x
48}
49
50macro_rules! impl_unsigned {
51    ($t: ident) => {
52        impl IsSquare for $t {
53            /// Determines whether an integer is a perfect square.
54            ///
55            /// $f(x) = (\exists b \in \Z : b^2 = x)$.
56            ///
57            /// # Worst-case complexity
58            /// Constant time and additional memory.
59            ///
60            /// # Examples
61            /// See [here](super::is_square#is_square).
62            #[inline]
63            fn is_square(&self) -> bool {
64                is_square_u64(u64::wrapping_from(*self))
65            }
66        }
67    };
68}
69impl_unsigned!(u8);
70impl_unsigned!(u16);
71impl_unsigned!(u32);
72impl_unsigned!(u64);
73impl_unsigned!(usize);
74
75// From mpn/generic/mod_34lsub1.c
76const B1: u64 = u64::WIDTH >> 2;
77const B2: u64 = B1 * 2;
78const B3: u64 = B1 * 3;
79
80const M2: u64 = (1 << B2) - 1;
81const M3: u64 = (1 << B3) - 1;
82
83const fn low0(n: u64) -> u64 {
84    n & M3
85}
86
87const fn high0(n: u64) -> u64 {
88    n >> B3
89}
90
91const fn low1(n: u64) -> u64 {
92    (n & M2) << B1
93}
94
95const fn high1(n: u64) -> u64 {
96    n >> B2
97}
98
99// This is mpn_mod_34lsub1 from mpn/generic/mod_34lsub1.c, GMP 6.3.0.
100//
101// Calculate a remainder from `limbs` divided by 2^(u64::WIDTH*3/4)-1. The remainder is not fully
102// reduced, it's any limb value congruent to `limbs` modulo that divisor.
103//
104// Check gen-psqr.c. mpn_mod_34lsub1 preferred over mpn_mod_1 (plus a PERFSQR_PP modulus) with 32
105// and 64 bit limb.
106const fn mod_34lsub1(x_hi: u64, x_lo: u64) -> u64 {
107    low0(x_lo) + high0(x_lo) + low1(x_hi) + high1(x_hi)
108}
109
110const MOD34_BITS: u64 = (u64::WIDTH >> 2) * 3;
111const MOD34_MASK: u64 = (1 << MOD34_BITS) - 1;
112
113const fn perfsqr_mod_34(x_hi: u64, x_lo: u64) -> u64 {
114    let r = mod_34lsub1(x_hi, x_lo);
115    (r & MOD34_MASK) + (r >> MOD34_BITS)
116}
117
118// This is PERFSQR_MOD_BITS from mpn/perfsqr.h, GMP 6.3.0. Either 49 on 64 bit limb or 25 on 32 bit
119// limb. 2^48-1 = 3^2 * 5 * 7 * 13 * 17 * 97 ... 2^24-1 = 3^2 * 5 * 7 * 13 * 17 ...
120const SQR_MOD_BITS: u64 = MOD34_BITS + 1;
121const SQR_MOD_MASK: u64 = (1 << SQR_MOD_BITS) - 1;
122
123const fn perfsqr_mod_idx(r: u64, d: u64, inv: u64) -> u64 {
124    assert!(r <= SQR_MOD_MASK);
125    assert!(inv.wrapping_mul(d) & SQR_MOD_MASK == 1);
126    assert!(u64::MAX / d >= SQR_MOD_MASK);
127    let q = r.wrapping_mul(inv) & SQR_MOD_MASK;
128    assert!(r == (q.wrapping_mul(d) & SQR_MOD_MASK));
129    q.wrapping_mul(d) >> SQR_MOD_BITS
130}
131
132// Single limb. Check precomputed bitmasks to see if remainder is a quadratic residue
133fn perfsqr_mod_1(r: u64, d: u64, inv: u64, mask: u64) -> bool {
134    assert!(d <= u64::WIDTH);
135    let idx = perfsqr_mod_idx(r, d, inv);
136    if (mask >> idx) & 1 == 0 {
137        // non-square
138        return false;
139    }
140    true
141}
142
143// Double limb. Check precomputed bitmasks to see if remainder is a quadratic residue
144fn perfsqr_mod_2(r: u64, d: u64, inv: u64, mhi: u64, mlo: u64) -> bool {
145    assert!(d <= const { 2 * u64::WIDTH });
146    let mut idx = perfsqr_mod_idx(r, d, inv);
147    let m = if idx < u64::WIDTH { mlo } else { mhi };
148    idx %= u64::WIDTH;
149    if (m >> idx) & 1 == 0 {
150        // non-square
151        return false;
152    }
153    true
154}
155
156// This test identifies 97.81% as non-squares. Grand total sq_res_0x100 and PERFSQR_MOD_TEST, 99.62%
157// non-squares.
158fn perfsqr_mod_test(x: u128) -> bool {
159    let (x_hi, x_lo) = x.split_in_half();
160    let r = perfsqr_mod_34(x_hi, x_lo);
161    perfsqr_mod_2(r, 91, 0xfd2fd2fd2fd3, 0x2191240, 0x8850a206953820e1) // 69.23%
162        && perfsqr_mod_2(r, 85, 0xfcfcfcfcfcfd, 0x82158, 0x10b48c4b4206a105) // 68.24%
163        && perfsqr_mod_1(r, 9, 0xe38e38e38e39, 0x93)  // 55.56%
164        && perfsqr_mod_2(r, 97, 0xfd5c5f02a3a1, 0x1eb628b47, 0x6067981b8b451b5f) // 49.48%
165}
166
167// This is sq_res0x100 from mpn/perfsqr.h when generated for 64 bit limb, GMP 6.3.0. Non-zero bit
168// indicates a quadratic residue mod 0x100. This test identifies 82.81% as non-squares (212/256).
169const SQR_MOD256: [u64; 4] =
170    [0x202021202030213, 0x202021202020213, 0x202021202030212, 0x202021202020212];
171
172impl IsSquare for u128 {
173    /// Determines whether an integer is a perfect square.
174    ///
175    /// $f(x) = (\exists b \in \Z : b^2 = x)$.
176    ///
177    /// # Worst-case complexity
178    /// Constant time and additional memory.
179    ///
180    /// # Examples
181    /// See [here](super::is_square#is_square).
182    #[inline]
183    fn is_square(&self) -> bool {
184        let idx = self % 0x100; // mod 256
185
186        // The first test excludes 212/256 (82.8%) of the perfect square candidates in O(1) time.
187        //
188        // This just checks the particular bit in the bitmask SQR_MOD256_U64 encoding where the
189        // input can be a perfect square mod 256.
190        if (SQR_MOD256[(idx >> u64::LOG_WIDTH) as usize]
191            >> (idx & const { u64::WIDTH_MASK as Self }))
192            & 1
193            == 0
194        {
195            return false;
196        }
197        // The second test uses mpn_mod_34lsub1 to detect non-squares according to their residues
198        // modulo small primes (or powers of primes). See mpn/perfsqr.h, GMP 6.3.0.
199        if !perfsqr_mod_test(*self) {
200            return false;
201        }
202        // For the third and last test, we finally compute the square root, to make sure we've
203        // really got a perfect square.
204        self.checked_sqrt().is_some()
205    }
206}
207
208macro_rules! impl_signed {
209    ($t: ident) => {
210        impl IsSquare for $t {
211            /// Determines whether an integer is a perfect square.
212            ///
213            /// $f(x) = (\exists b \in \Z : b^2 = x)$.
214            ///
215            /// # Worst-case complexity
216            /// Constant time and additional memory.
217            ///
218            /// # Examples
219            /// See [here](super::is_square#is_square).
220            #[inline]
221            fn is_square(&self) -> bool {
222                if *self < 0 {
223                    false
224                } else {
225                    self.unsigned_abs().is_square()
226                }
227            }
228        }
229    };
230}
231apply_to_signeds!(impl_signed);
malachite_base/num/factorization/is_square.rs

malachite_base/num/factorization/
is_square.rs
