num_prime/
integer.rs

1//! Backend implementations for integers
2
3#[cfg(feature = "num-bigint")]
4use crate::tables::{CUBIC_MODULI, CUBIC_RESIDUAL, QUAD_MODULI, QUAD_RESIDUAL};
5use crate::traits::{BitTest, ExactRoots};
6
7use num_integer::Roots;
8
9#[cfg(feature = "num-bigint")]
10use num_bigint::{BigInt, BigUint, ToBigInt};
11#[cfg(feature = "num-bigint")]
12use num_traits::{One, Signed, ToPrimitive, Zero};
13
14macro_rules! impl_bittest_prim {
15    ($($T:ty)*) => {$(
16        impl BitTest for $T {
17            #[inline]
18            fn bits(&self) -> usize {
19                (<$T>::BITS - self.leading_zeros()) as usize
20            }
21            #[inline]
22            fn bit(&self, position: usize) -> bool {
23                self & (1 << position) > 0
24            }
25            #[inline]
26            fn trailing_zeros(&self) -> usize {
27                <$T>::trailing_zeros(*self) as usize
28            }
29        }
30    )*}
31}
32impl_bittest_prim!(u8 u16 u32 u64 u128 usize);
33
34#[cfg(feature = "num-bigint")]
35impl BitTest for BigUint {
36    fn bit(&self, position: usize) -> bool {
37        self.bit(position as u64)
38    }
39    fn bits(&self) -> usize {
40        BigUint::bits(self) as usize
41    }
42    #[inline]
43    fn trailing_zeros(&self) -> usize {
44        match BigUint::trailing_zeros(self) {
45            Some(a) => a as usize,
46            None => 0,
47        }
48    }
49}
50
51macro_rules! impl_exactroot_prim {
52    ($($T:ty)*) => {$(
53        impl ExactRoots for $T {
54            fn nth_root_exact(&self, n: u32) -> Option<Self> {
55                // For even roots of negative numbers, return None instead of panicking
56                if self < &0 && n % 2 == 0 {
57                    return None;
58                }
59                let r = self.nth_root(n);
60                if &r.clone().pow(n) == self {
61                    Some(r)
62                } else {
63                    None
64                }
65            }
66            fn sqrt_exact(&self) -> Option<Self> {
67                if self < &0 { return None; }
68                let shift = self.trailing_zeros();
69
70                // the general form of any square number is (2^(2m))(8N+1)
71                if shift & 1 == 1 { return None; }
72                if (self >> shift) & 7 != 1 { return None; }
73                self.nth_root_exact(2)
74            }
75        }
76    )*};
77    // TODO: it might worth use QUAD_RESIDUE and CUBIC_RESIDUE for large size
78    //       primitive integers, need benchmark
79}
80impl_exactroot_prim!(u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
81
82#[cfg(feature = "num-bigint")]
83impl ExactRoots for BigUint {
84    fn sqrt_exact(&self) -> Option<Self> {
85        // shortcuts
86        if self.is_zero() {
87            return Some(BigUint::zero());
88        }
89        if let Some(v) = self.to_u64() {
90            return v.sqrt_exact().map(BigUint::from);
91        }
92
93        // check mod 2
94        let shift = self.trailing_zeros().unwrap();
95        if shift & 1 == 1 {
96            return None;
97        }
98        if !((self >> shift) & BigUint::from(7u8)).is_one() {
99            return None;
100        }
101
102        // check other moduli
103        #[cfg(not(feature = "big-table"))]
104        for (m, res) in QUAD_MODULI.iter().zip(QUAD_RESIDUAL) {
105            // need to &63 since we have 65 in QUAD_MODULI
106            if (res >> ((self % m).to_u8().unwrap() & 63)) & 1 == 0 {
107                return None;
108            }
109        }
110        #[cfg(feature = "big-table")]
111        for (m, res) in QUAD_MODULI.iter().zip(QUAD_RESIDUAL) {
112            let rem = (self % m).to_u16().unwrap();
113            if (res[(rem / 64) as usize] >> (rem % 64)) & 1 == 0 {
114                return None;
115            }
116        }
117
118        self.nth_root_exact(2)
119    }
120
121    fn cbrt_exact(&self) -> Option<Self> {
122        // shortcuts
123        if self.is_zero() {
124            return Some(BigUint::zero());
125        }
126        if let Some(v) = self.to_u64() {
127            return v.cbrt_exact().map(BigUint::from);
128        }
129
130        // check mod 2
131        let shift = self.trailing_zeros().unwrap();
132        if shift % 3 != 0 {
133            return None;
134        }
135
136        // check other moduli
137        #[cfg(not(feature = "big-table"))]
138        for (m, res) in CUBIC_MODULI.iter().zip(CUBIC_RESIDUAL) {
139            if (res >> (self % m).to_u8().unwrap()) & 1 == 0 {
140                return None;
141            }
142        }
143        #[cfg(feature = "big-table")]
144        for (m, res) in CUBIC_MODULI.iter().zip(CUBIC_RESIDUAL) {
145            let rem = (self % m).to_u16().unwrap();
146            if (res[(rem / 64) as usize] >> (rem % 64)) & 1 == 0 {
147                return None;
148            }
149        }
150
151        self.nth_root_exact(3)
152    }
153}
154
155#[cfg(feature = "num-bigint")]
156impl ExactRoots for BigInt {
157    fn nth_root_exact(&self, n: u32) -> Option<Self> {
158        // For even roots of negative numbers, return None instead of panicking
159        if self.is_negative() && n % 2 == 0 {
160            return None;
161        }
162
163        // For odd roots, handle negative numbers by taking the root of the magnitude
164        // and then applying the sign
165        if self.is_negative() {
166            self.magnitude()
167                .nth_root_exact(n)
168                .and_then(|u| u.to_bigint())
169                .map(|v| -v)
170        } else {
171            self.magnitude()
172                .nth_root_exact(n)
173                .and_then(|u| u.to_bigint())
174        }
175    }
176    fn sqrt_exact(&self) -> Option<Self> {
177        self.to_biguint()
178            .and_then(|u| u.sqrt_exact())
179            .and_then(|u| u.to_bigint())
180    }
181    fn cbrt_exact(&self) -> Option<Self> {
182        self.magnitude()
183            .cbrt_exact()
184            .and_then(|u| u.to_bigint())
185            .map(|v| if self.is_negative() { -v } else { v })
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn exact_root_test() {
195        // some simple tests
196        assert!(ExactRoots::sqrt_exact(&3u8).is_none());
197        assert!(matches!(ExactRoots::sqrt_exact(&4u8), Some(2)));
198        assert!(matches!(ExactRoots::sqrt_exact(&9u8), Some(3)));
199        assert!(ExactRoots::sqrt_exact(&18u8).is_none());
200        assert!(ExactRoots::sqrt_exact(&3i8).is_none());
201        assert!(matches!(ExactRoots::sqrt_exact(&4i8), Some(2)));
202        assert!(matches!(ExactRoots::sqrt_exact(&9i8), Some(3)));
203        assert!(ExactRoots::sqrt_exact(&18i8).is_none());
204
205        // test fast implementations of sqrt against nth_root
206        for _ in 0..100 {
207            let x = rand::random::<u32>();
208            assert_eq!(
209                ExactRoots::sqrt_exact(&x),
210                ExactRoots::nth_root_exact(&x, 2)
211            );
212            assert_eq!(
213                ExactRoots::cbrt_exact(&x),
214                ExactRoots::nth_root_exact(&x, 3)
215            );
216            let x = rand::random::<i32>();
217            assert_eq!(
218                ExactRoots::cbrt_exact(&x),
219                ExactRoots::nth_root_exact(&x, 3)
220            );
221        }
222        // test perfect powers
223        for _ in 0..100 {
224            let x = u64::from(rand::random::<u32>());
225            assert!(matches!(ExactRoots::sqrt_exact(&(x * x)), Some(v) if v == x));
226            let x = i64::from(rand::random::<i16>());
227            assert!(matches!(ExactRoots::cbrt_exact(&(x * x * x)), Some(v) if v == x));
228        }
229        // test non-perfect powers
230        for _ in 0..100 {
231            let x = u64::from(rand::random::<u32>());
232            let y = u64::from(rand::random::<u32>());
233            if x == y {
234                continue;
235            }
236            assert!(ExactRoots::sqrt_exact(&(x * y)).is_none());
237        }
238
239        #[cfg(feature = "num-bigint")]
240        {
241            use num_bigint::RandBigInt;
242            let mut rng = rand::thread_rng();
243            // test fast implementations of sqrt against nth_root
244            for _ in 0..10 {
245                let x = rng.gen_biguint(150);
246                assert_eq!(
247                    ExactRoots::sqrt_exact(&x),
248                    ExactRoots::nth_root_exact(&x, 2)
249                );
250                assert_eq!(
251                    ExactRoots::cbrt_exact(&x),
252                    ExactRoots::nth_root_exact(&x, 3)
253                );
254                let x = rng.gen_bigint(150);
255                assert_eq!(
256                    ExactRoots::cbrt_exact(&x),
257                    ExactRoots::nth_root_exact(&x, 3)
258                );
259            }
260            // test perfect powers
261            for _ in 0..10 {
262                let x = rng.gen_biguint(150);
263                assert!(matches!(ExactRoots::sqrt_exact(&(&x * &x)), Some(v) if v == x));
264                let x = rng.gen_biguint(150);
265                assert!(
266                    matches!(ExactRoots::cbrt_exact(&(&x * &x * &x)), Some(v) if v == x),
267                    "failed at {}",
268                    x
269                );
270            }
271            // test non-perfect powers
272            for _ in 0..10 {
273                let x = rng.gen_biguint(150);
274                let y = rng.gen_biguint(150);
275                if x == y {
276                    continue;
277                }
278                assert!(ExactRoots::sqrt_exact(&(x * y)).is_none());
279            }
280        }
281    }
282
283    #[test]
284    fn test_nth_root_exact_negative_even_root() {
285        // Test for issue #25: nth_root_exact(2) should return None for negative numbers
286        // instead of panicking
287        let result = (-1i32).nth_root_exact(2);
288        assert!(
289            result.is_none(),
290            "nth_root_exact(2) should return None for negative numbers"
291        );
292
293        // Test other negative numbers with even roots
294        let result = (-4i32).nth_root_exact(2);
295        assert!(
296            result.is_none(),
297            "nth_root_exact(2) should return None for negative numbers"
298        );
299
300        let result = (-8i32).nth_root_exact(4);
301        assert!(
302            result.is_none(),
303            "nth_root_exact(4) should return None for negative numbers"
304        );
305
306        // Test that odd roots of negative numbers still work
307        let result = (-8i32).nth_root_exact(3);
308        assert_eq!(
309            result,
310            Some(-2),
311            "nth_root_exact(3) should work for negative numbers"
312        );
313
314        let result = (-27i32).nth_root_exact(3);
315        assert_eq!(
316            result,
317            Some(-3),
318            "nth_root_exact(3) should work for negative numbers"
319        );
320    }
321
322    #[test]
323    fn test_nth_root_exact_all_signed_types() {
324        // Test all signed integer types with even roots of negative numbers
325        assert_eq!((-1i8).nth_root_exact(2), None);
326        assert_eq!((-1i16).nth_root_exact(2), None);
327        assert_eq!((-1i32).nth_root_exact(2), None);
328        assert_eq!((-1i64).nth_root_exact(2), None);
329        assert_eq!((-1i128).nth_root_exact(2), None);
330        assert_eq!((-1isize).nth_root_exact(2), None);
331
332        // Test odd roots work correctly
333        assert_eq!((-8i32).nth_root_exact(3), Some(-2));
334        assert_eq!((-32i32).nth_root_exact(5), Some(-2));
335
336        // Test positive cases still work
337        assert_eq!(16i32.nth_root_exact(4), Some(2));
338        assert_eq!(32i32.nth_root_exact(5), Some(2));
339    }
340
341    #[test]
342    #[cfg(feature = "num-bigint")]
343    fn test_nth_root_exact_bigint_negative() {
344        use num_bigint::BigInt;
345
346        // Test even roots return None for negative BigInt
347        assert_eq!(BigInt::from(-1).nth_root_exact(2), None);
348        assert_eq!(BigInt::from(-16).nth_root_exact(4), None);
349
350        // Test odd roots work for negative BigInt
351        assert_eq!(BigInt::from(-8).nth_root_exact(3), Some(BigInt::from(-2)));
352        assert_eq!(
353            BigInt::from(-1000000000i64).nth_root_exact(3),
354            Some(BigInt::from(-1000i32))
355        );
356    }
357}
num_prime/integer.rs

num_prime/
integer.rs
