num_prime/
integer.rs

1//! Backend implementations for integers
2
3use crate::tables::{CUBIC_MODULI, CUBIC_RESIDUAL, QUAD_MODULI, QUAD_RESIDUAL};
4use crate::traits::{BitTest, ExactRoots};
5
6#[cfg(feature = "num-bigint")]
7use num_bigint::{BigInt, BigUint, ToBigInt};
8#[cfg(feature = "num-bigint")]
9use num_traits::{One, Signed, ToPrimitive, Zero};
10
11macro_rules! impl_bittest_prim {
12    ($($T:ty)*) => {$(
13        impl BitTest for $T {
14            #[inline]
15            fn bits(&self) -> usize {
16                (<$T>::BITS - self.leading_zeros()) as usize
17            }
18            #[inline]
19            fn bit(&self, position: usize) -> bool {
20                self & (1 << position) > 0
21            }
22            #[inline]
23            fn trailing_zeros(&self) -> usize {
24                <$T>::trailing_zeros(*self) as usize
25            }
26        }
27    )*}
28}
29impl_bittest_prim!(u8 u16 u32 u64 u128 usize);
30
31#[cfg(feature = "num-bigint")]
32impl BitTest for BigUint {
33    fn bit(&self, position: usize) -> bool {
34        self.bit(position as u64)
35    }
36    fn bits(&self) -> usize {
37        BigUint::bits(&self) as usize
38    }
39    #[inline]
40    fn trailing_zeros(&self) -> usize {
41        match BigUint::trailing_zeros(&self) {
42            Some(a) => a as usize,
43            None => 0,
44        }
45    }
46}
47
48macro_rules! impl_exactroot_prim {
49    ($($T:ty)*) => {$(
50        impl ExactRoots for $T {
51            fn sqrt_exact(&self) -> Option<Self> {
52                if self < &0 { return None; }
53                let shift = self.trailing_zeros();
54
55                // the general form of any square number is (2^(2m))(8N+1)
56                if shift & 1 == 1 { return None; }
57                if (self >> shift) & 7 != 1 { return None; }
58                self.nth_root_exact(2)
59            }
60        }
61    )*};
62    // TODO: it might worth use QUAD_RESIDUE and CUBIC_RESIDUE for large size
63    //       primitive integers, need benchmark
64}
65impl_exactroot_prim!(u8 u16 u32 u64 u128 usize i8 i16 i32 i64 i128 isize);
66
67#[cfg(feature = "num-bigint")]
68impl ExactRoots for BigUint {
69    fn sqrt_exact(&self) -> Option<Self> {
70        // shortcuts
71        if self.is_zero() {
72            return Some(BigUint::zero());
73        }
74        if let Some(v) = self.to_u64() {
75            return v.sqrt_exact().map(BigUint::from);
76        }
77
78        // check mod 2
79        let shift = self.trailing_zeros().unwrap();
80        if shift & 1 == 1 {
81            return None;
82        }
83        if !((self >> shift) & BigUint::from(7u8)).is_one() {
84            return None;
85        }
86
87        // check other moduli
88        #[cfg(not(feature = "big-table"))]
89        for (m, res) in QUAD_MODULI.iter().zip(QUAD_RESIDUAL) {
90            // need to &63 since we have 65 in QUAD_MODULI
91            if (res >> ((self % m).to_u8().unwrap() & 63)) & 1 == 0 {
92                return None;
93            }
94        }
95        #[cfg(feature = "big-table")]
96        for (m, res) in QUAD_MODULI.iter().zip(QUAD_RESIDUAL) {
97            let rem = (self % m).to_u16().unwrap();
98            if (res[(rem / 64) as usize] >> (rem % 64)) & 1 == 0 {
99                return None;
100            }
101        }
102
103        self.nth_root_exact(2)
104    }
105
106    fn cbrt_exact(&self) -> Option<Self> {
107        // shortcuts
108        if self.is_zero() {
109            return Some(BigUint::zero());
110        }
111        if let Some(v) = self.to_u64() {
112            return v.cbrt_exact().map(BigUint::from);
113        }
114
115        // check mod 2
116        let shift = self.trailing_zeros().unwrap();
117        if shift % 3 != 0 {
118            return None;
119        }
120
121        // check other moduli
122        #[cfg(not(feature = "big-table"))]
123        for (m, res) in CUBIC_MODULI.iter().zip(CUBIC_RESIDUAL) {
124            if (res >> (self % m).to_u8().unwrap()) & 1 == 0 {
125                return None;
126            }
127        }
128        #[cfg(feature = "big-table")]
129        for (m, res) in CUBIC_MODULI.iter().zip(CUBIC_RESIDUAL) {
130            let rem = (self % m).to_u16().unwrap();
131            if (res[(rem / 64) as usize] >> (rem % 64)) & 1 == 0 {
132                return None;
133            }
134        }
135
136        self.nth_root_exact(3)
137    }
138}
139
140#[cfg(feature = "num-bigint")]
141impl ExactRoots for BigInt {
142    fn sqrt_exact(&self) -> Option<Self> {
143        self.to_biguint()
144            .and_then(|u| u.sqrt_exact())
145            .and_then(|u| u.to_bigint())
146    }
147    fn cbrt_exact(&self) -> Option<Self> {
148        self.magnitude()
149            .cbrt_exact()
150            .and_then(|u| u.to_bigint())
151            .map(|v| if self.is_negative() { -v } else { v })
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use rand;
159
160    #[test]
161    fn exact_root_test() {
162        // some simple tests
163        assert!(matches!(ExactRoots::sqrt_exact(&3u8), None));
164        assert!(matches!(ExactRoots::sqrt_exact(&4u8), Some(2)));
165        assert!(matches!(ExactRoots::sqrt_exact(&9u8), Some(3)));
166        assert!(matches!(ExactRoots::sqrt_exact(&18u8), None));
167        assert!(matches!(ExactRoots::sqrt_exact(&3i8), None));
168        assert!(matches!(ExactRoots::sqrt_exact(&4i8), Some(2)));
169        assert!(matches!(ExactRoots::sqrt_exact(&9i8), Some(3)));
170        assert!(matches!(ExactRoots::sqrt_exact(&18i8), None));
171
172        // test fast implementations of sqrt against nth_root
173        for _ in 0..100 {
174            let x = rand::random::<u32>();
175            assert_eq!(
176                ExactRoots::sqrt_exact(&x),
177                ExactRoots::nth_root_exact(&x, 2)
178            );
179            assert_eq!(
180                ExactRoots::cbrt_exact(&x),
181                ExactRoots::nth_root_exact(&x, 3)
182            );
183            let x = rand::random::<i32>();
184            assert_eq!(
185                ExactRoots::cbrt_exact(&x),
186                ExactRoots::nth_root_exact(&x, 3)
187            );
188        }
189        // test perfect powers
190        for _ in 0..100 {
191            let x = rand::random::<u32>() as u64;
192            assert!(matches!(ExactRoots::sqrt_exact(&(x * x)), Some(v) if v == x));
193            let x = rand::random::<i16>() as i64;
194            assert!(matches!(ExactRoots::cbrt_exact(&(x * x * x)), Some(v) if v == x));
195        }
196        // test non-perfect powers
197        for _ in 0..100 {
198            let x = rand::random::<u32>() as u64;
199            let y = rand::random::<u32>() as u64;
200            if x == y {
201                continue;
202            }
203            assert!(ExactRoots::sqrt_exact(&(x * y)).is_none());
204        }
205
206        #[cfg(feature = "num-bigint")]
207        {
208            use num_bigint::RandBigInt;
209            let mut rng = rand::thread_rng();
210            // test fast implementations of sqrt against nth_root
211            for _ in 0..10 {
212                let x = rng.gen_biguint(150);
213                assert_eq!(
214                    ExactRoots::sqrt_exact(&x),
215                    ExactRoots::nth_root_exact(&x, 2)
216                );
217                assert_eq!(
218                    ExactRoots::cbrt_exact(&x),
219                    ExactRoots::nth_root_exact(&x, 3)
220                );
221                let x = rng.gen_bigint(150);
222                assert_eq!(
223                    ExactRoots::cbrt_exact(&x),
224                    ExactRoots::nth_root_exact(&x, 3)
225                );
226            }
227            // test perfect powers
228            for _ in 0..10 {
229                let x = rng.gen_biguint(150);
230                assert!(matches!(ExactRoots::sqrt_exact(&(&x * &x)), Some(v) if v == x));
231                let x = rng.gen_biguint(150);
232                assert!(
233                    matches!(ExactRoots::cbrt_exact(&(&x * &x * &x)), Some(v) if v == x),
234                    "failed at {}",
235                    x
236                );
237            }
238            // test non-perfect powers
239            for _ in 0..10 {
240                let x = rng.gen_biguint(150);
241                let y = rng.gen_biguint(150);
242                if x == y {
243                    continue;
244                }
245                assert!(ExactRoots::sqrt_exact(&(x * y)).is_none());
246            }
247        }
248    }
249}
num_prime/integer.rs

num_prime/
integer.rs
