Skip to main content

num_prime/
traits.rs

1use core::default::Default;
2use std::ops::{BitAnd, BitOr, Mul};
3
4use either::Either;
5use num_integer::{Integer, Roots};
6use num_traits::Pow;
7
8/// This trait support unified bit testing for (unsigned) integers
9pub trait BitTest {
10    /// Get the minimum required number of bits to represent this integer
11    fn bits(&self) -> usize;
12
13    /// Get the i-th bit of the integer, with i specified by `position`
14    fn bit(&self, position: usize) -> bool;
15
16    /// Get the number of trailing zeros in the integer
17    fn trailing_zeros(&self) -> usize;
18}
19
20/// This enum describes the result of primality checks
21#[derive(Debug, Clone, Copy, PartialEq)]
22pub enum Primality {
23    /// The number passes deterministic primality check.
24    Yes,
25    /// The number is a composite and failed at least one specific primality check.
26    No,
27    /// The number passes several probabilistic primality check.
28    /// The associated float number carries the probability of the number being a prime
29    /// (conditioned on that it's already a probable prime)
30    Probable(f32),
31}
32
33impl Primality {
34    /// Check whether the resule indicates that the number is
35    /// (very) probably a prime. Return false only on [`Primality::No`]
36    #[inline(always)]
37    #[must_use]
38    pub fn probably(self) -> bool {
39        !matches!(self, Primality::No)
40    }
41}
42
43impl BitAnd<Primality> for Primality {
44    type Output = Primality;
45
46    /// Combine two primality results by ensuring both numbers are prime
47    fn bitand(self, rhs: Primality) -> Self::Output {
48        match self {
49            Primality::No => Primality::No,
50            Primality::Yes => rhs,
51            Primality::Probable(p) => match rhs {
52                Primality::No => Primality::No,
53                Primality::Yes => Primality::Probable(p),
54                Primality::Probable(p2) => {
55                    let combined = p.mul(p2);
56                    Primality::Probable(combined)
57                }
58            },
59        }
60    }
61}
62
63impl BitOr<Primality> for Primality {
64    type Output = Primality;
65
66    /// Combine two primality results by ensuring either numbers is prime
67    fn bitor(self, rhs: Primality) -> Self::Output {
68        match self {
69            Primality::No => rhs,
70            Primality::Yes => Primality::Yes,
71            Primality::Probable(p) => match rhs {
72                Primality::No => Primality::Probable(p),
73                Primality::Yes => Primality::Yes,
74                Primality::Probable(p2) => Primality::Probable(1. - (1. - p) * (1. - p2)),
75            },
76        }
77    }
78}
79
80/// Represents a configuration for a primality test
81#[derive(Debug, Clone, Copy)]
82#[non_exhaustive]
83pub struct PrimalityTestConfig {
84    // TODO: add option to divides small primes in the table
85    //       and this option should be enabled if the probabilistic test is used for strict config
86    /// Number of strong probable prime test, starting from base 2
87    pub sprp_trials: usize,
88
89    /// Number of strong probable prime test with random bases
90    pub sprp_random_trials: usize,
91
92    /// Whether perform strong lucas probable prime test (with automatically selected parameters)
93    pub slprp_test: bool,
94
95    /// Whether perform extra strong lucas probable prime test (with automatically selected parameters)
96    pub eslprp_test: bool,
97}
98
99impl Default for PrimalityTestConfig {
100    /// Create a defalt primality testing configuration. This config will eliminate most
101    /// composites with little computation
102    fn default() -> Self {
103        Self {
104            sprp_trials: 2,        // test base 2 and 3
105            sprp_random_trials: 3, // choose other 3 random bases
106            slprp_test: false,
107            eslprp_test: false,
108        }
109    }
110}
111
112impl PrimalityTestConfig {
113    /// Create a configuration with a very strong primality check. It's based on
114    /// the **strongest deterministic primality testing** and some SPRP tests with
115    /// random bases.
116    #[must_use]
117    pub fn strict() -> Self {
118        let mut config = Self::bpsw();
119        config.sprp_random_trials = 1;
120        config
121    }
122
123    /// Create a configuration for Baillie-PSW test (base 2 SPRP test + SLPRP test)
124    #[must_use]
125    pub fn bpsw() -> Self {
126        Self {
127            sprp_trials: 1,
128            sprp_random_trials: 0,
129            slprp_test: true,
130            eslprp_test: false,
131        }
132    }
133}
134
135/// Represents a configuration for integer factorization
136#[derive(Debug, Clone, Copy)]
137#[non_exhaustive]
138pub struct FactorizationConfig {
139    /// Config for testing if a factor is prime
140    pub primality_config: PrimalityTestConfig,
141
142    /// Prime limit of trial division, you also need to reserve the primes in the buffer
143    /// if all primes under the limit are to be tested. `None` means using all available primes.
144    pub td_limit: Option<u64>,
145
146    /// Number of trials with Pollard's rho method
147    pub rho_trials: usize,
148}
149
150impl Default for FactorizationConfig {
151    /// Create a defalt primality testing configuration. This config will factorize
152    /// most integers within decent time
153    fn default() -> Self {
154        const THRESHOLD_DEFAULT_TD: u64 = 1 << 14;
155        Self {
156            primality_config: PrimalityTestConfig::default(),
157            td_limit: Some(THRESHOLD_DEFAULT_TD),
158            rho_trials: 4,
159        }
160    }
161}
162
163impl FactorizationConfig {
164    /// Same as the default configuration but with strict primality check
165    #[must_use]
166    pub fn strict() -> Self {
167        Self {
168            primality_config: PrimalityTestConfig::strict(),
169            ..Self::default()
170        }
171    }
172}
173
174// FIXME: backport to num_integer (see https://github.com/rust-num/num-traits/issues/233)
175/// Extension on [`num_integer::Roots`] to support perfect power check on integers
176pub trait ExactRoots: Roots + Pow<u32, Output = Self> + Clone {
177    fn nth_root_exact(&self, n: u32) -> Option<Self> {
178        let r = self.nth_root(n);
179        if &r.clone().pow(n) == self {
180            Some(r)
181        } else {
182            None
183        }
184    }
185    fn sqrt_exact(&self) -> Option<Self> {
186        self.nth_root_exact(2)
187    }
188    fn cbrt_exact(&self) -> Option<Self> {
189        self.nth_root_exact(3)
190    }
191    fn is_nth_power(&self, n: u32) -> bool {
192        self.nth_root_exact(n).is_some()
193    }
194    fn is_square(&self) -> bool {
195        self.sqrt_exact().is_some()
196    }
197    fn is_cubic(&self) -> bool {
198        self.cbrt_exact().is_some()
199    }
200}
201
202// TODO: implement is_perfect_power, this could be used during factorization
203//       to filter out perfect powers when factorize large integers
204// REF: PARI/GP `Z_ispowerall`, `is_357_power`
205//      FLINT `n_is_perfect_power235`, `fmpz_is_perfect_power`
206//      GMP `mpz_perfect_power_p`
207
208/// This trait represents a general data structure that stores primes.
209///
210/// It's recommended to store at least a bunch of small primes in the buffer
211/// to make some of the algorithms more efficient.
212pub trait PrimeBuffer<'a> {
213    type PrimeIter: Iterator<Item = &'a u64>;
214
215    /// Directly return an iterator of existing primes
216    fn iter(&'a self) -> Self::PrimeIter;
217
218    /// Generate primes until the largest prime in the buffer is equal or larger than limit
219    fn reserve(&mut self, limit: u64);
220
221    /// Get the largest primes in the list
222    fn bound(&self) -> u64;
223
224    /// Test if the number is in the buffer. If a number is not in the buffer,
225    /// then it's either a composite or large than [`PrimeBuffer::bound()`]
226    fn contains(&self, num: u64) -> bool;
227
228    /// clear the prime buffer to save memory
229    fn clear(&mut self);
230}
231
232/// This trait implements various primality testing algorithms
233///
234/// Reference:
235/// - <http://ntheory.org/pseudoprimes.html>
236pub trait PrimalityUtils: Integer + Clone {
237    /// Test if the integer is a (Fermat) probable prime
238    fn is_prp(&self, base: Self) -> bool;
239
240    /// Test if the integer is a strong probable prime (based on Miller-Rabin test).
241    fn is_sprp(&self, base: Self) -> bool;
242
243    /// Do a Miller-Rabin test. The return value is a integer if it finds a factor of
244    /// the integer, otherwise it reports the test result.
245    fn test_sprp(&self, base: Self) -> Either<bool, Self>;
246
247    /// Test if the integer is a Lucas probable prime
248    /// If either of p, q is not specified, then we will use Selfridge's Method A to choose p, q
249    fn is_lprp(&self, p: Option<usize>, q: Option<isize>) -> bool;
250
251    /// Test if the integer is a strong Lucas probable prime
252    /// If either of p, q is not specified, then we will use Selfridge's Method A to choose p, q
253    fn is_slprp(&self, p: Option<usize>, q: Option<isize>) -> bool;
254
255    /// Test if the integer is an extra strong Lucas probable prime
256    /// If p is not specified, then first p starting from 3 such that Jacobi symbol is -1 will be chosen, which is sometimes refered as "Method C"
257    fn is_eslprp(&self, p: Option<usize>) -> bool;
258
259    // TODO: implement ECPP test
260    // https://en.wikipedia.org/wiki/Elliptic_curve_primality
261
262    // TODO: implement is_vprp (Lucas-V probable prime test)
263    // https://arxiv.org/pdf/2006.14425.pdf
264}
265
266/// Supports random generation of primes
267pub trait RandPrime<T> {
268    /// Generate a random prime within the given bit size limit
269    ///
270    /// # Panics
271    /// if the `bit_size` is 0 or it's larger than the bit width of the integer
272    fn gen_prime(&mut self, bit_size: usize, config: Option<PrimalityTestConfig>) -> T;
273
274    /// Generate a random prime with **exact** the given bit size
275    ///
276    /// # Panics
277    /// if the `bit_size` is 0 or it's larger than the bit width of the integer
278    fn gen_prime_exact(&mut self, bit_size: usize, config: Option<PrimalityTestConfig>) -> T;
279
280    /// Generate a random (Sophie German) safe prime within the given bit size limit. The generated prime
281    /// is guaranteed to pass the [`is_safe_prime`][crate::nt_funcs::is_safe_prime] test
282    ///
283    /// # Panics
284    /// if the `bit_size` is 0 or it's larger than the bit width of the integer
285    fn gen_safe_prime(&mut self, bit_size: usize) -> T;
286
287    /// Generate a random (Sophie German) safe prime with the **exact** given bit size. The generated prime
288    /// is guaranteed to pass the [`is_safe_prime`][crate::nt_funcs::is_safe_prime] test
289    ///
290    /// # Panics
291    /// if the `bit_size` is 0 or it's larger than the bit width of the integer
292    fn gen_safe_prime_exact(&mut self, bit_size: usize) -> T;
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    // --- Primality::probably() ---
300    #[test]
301    fn primality_probably() {
302        assert!(Primality::Yes.probably());
303        assert!(!Primality::No.probably());
304        assert!(Primality::Probable(0.9).probably());
305    }
306
307    // --- Primality BitAnd: all 9 combinations ---
308    #[test]
309    fn primality_bitand_no_lhs() {
310        assert_eq!(Primality::No & Primality::No, Primality::No);
311        assert_eq!(Primality::No & Primality::Yes, Primality::No);
312        assert_eq!(Primality::No & Primality::Probable(0.9), Primality::No);
313    }
314
315    #[test]
316    fn primality_bitand_yes_lhs() {
317        assert_eq!(Primality::Yes & Primality::No, Primality::No);
318        assert_eq!(Primality::Yes & Primality::Yes, Primality::Yes);
319        assert_eq!(
320            Primality::Yes & Primality::Probable(0.8),
321            Primality::Probable(0.8)
322        );
323    }
324
325    #[test]
326    fn primality_bitand_probable_lhs() {
327        assert_eq!(Primality::Probable(0.9) & Primality::No, Primality::No);
328        assert_eq!(
329            Primality::Probable(0.9) & Primality::Yes,
330            Primality::Probable(0.9)
331        );
332        // Probable & Probable => combined probability
333        let result = Primality::Probable(0.8) & Primality::Probable(0.9);
334        match result {
335            Primality::Probable(p) => assert!((p - 0.72).abs() < 1e-6),
336            _ => panic!("expected Probable"),
337        }
338    }
339
340    // --- Primality BitOr: all 9 combinations ---
341    #[test]
342    fn primality_bitor_no_lhs() {
343        assert_eq!(Primality::No | Primality::No, Primality::No);
344        assert_eq!(Primality::No | Primality::Yes, Primality::Yes);
345        assert_eq!(
346            Primality::No | Primality::Probable(0.8),
347            Primality::Probable(0.8)
348        );
349    }
350
351    #[test]
352    fn primality_bitor_yes_lhs() {
353        assert_eq!(Primality::Yes | Primality::No, Primality::Yes);
354        assert_eq!(Primality::Yes | Primality::Yes, Primality::Yes);
355        assert_eq!(Primality::Yes | Primality::Probable(0.8), Primality::Yes);
356    }
357
358    #[test]
359    fn primality_bitor_probable_lhs() {
360        assert_eq!(
361            Primality::Probable(0.9) | Primality::No,
362            Primality::Probable(0.9)
363        );
364        assert_eq!(Primality::Probable(0.9) | Primality::Yes, Primality::Yes);
365        // Probable | Probable => 1 - (1-p1)*(1-p2)
366        let result = Primality::Probable(0.8) | Primality::Probable(0.9);
367        match result {
368            Primality::Probable(p) => assert!((p - 0.98).abs() < 1e-6),
369            _ => panic!("expected Probable"),
370        }
371    }
372
373    // --- PrimalityTestConfig constructors ---
374    #[test]
375    fn primality_test_config_default() {
376        let c = PrimalityTestConfig::default();
377        assert_eq!(c.sprp_trials, 2);
378        assert_eq!(c.sprp_random_trials, 3);
379        assert!(!c.slprp_test);
380        assert!(!c.eslprp_test);
381    }
382
383    #[test]
384    fn primality_test_config_bpsw() {
385        let c = PrimalityTestConfig::bpsw();
386        assert_eq!(c.sprp_trials, 1);
387        assert_eq!(c.sprp_random_trials, 0);
388        assert!(c.slprp_test);
389        assert!(!c.eslprp_test);
390    }
391
392    #[test]
393    fn primality_test_config_strict() {
394        let c = PrimalityTestConfig::strict();
395        assert_eq!(c.sprp_trials, 1);
396        assert_eq!(c.sprp_random_trials, 1);
397        assert!(c.slprp_test);
398        assert!(!c.eslprp_test);
399    }
400
401    // --- FactorizationConfig constructors ---
402    #[test]
403    fn factorization_config_default() {
404        let c = FactorizationConfig::default();
405        assert_eq!(c.td_limit, Some(1 << 14));
406        assert_eq!(c.rho_trials, 4);
407    }
408
409    #[test]
410    fn factorization_config_strict() {
411        let c = FactorizationConfig::strict();
412        assert_eq!(c.td_limit, Some(1 << 14));
413        assert_eq!(c.rho_trials, 4);
414        // uses strict primality config
415        assert_eq!(c.primality_config.sprp_trials, 1);
416        assert_eq!(c.primality_config.sprp_random_trials, 1);
417        assert!(c.primality_config.slprp_test);
418    }
419
420    // --- ExactRoots helper methods ---
421    #[test]
422    fn exact_roots_is_square() {
423        assert!(16u64.is_square());
424        assert!(!15u64.is_square());
425        assert!(1u64.is_square());
426        assert!(100u64.is_square());
427        assert!(!99u64.is_square());
428    }
429
430    #[test]
431    fn exact_roots_is_cubic() {
432        assert!(27u64.is_cubic());
433        assert!(!26u64.is_cubic());
434        assert!(1u64.is_cubic());
435        assert!(64u64.is_cubic());
436    }
437
438    #[test]
439    fn exact_roots_is_nth_power() {
440        assert!(256u64.is_nth_power(4));
441        assert!(!255u64.is_nth_power(4));
442    }
443
444    #[test]
445    fn exact_roots_sqrt_exact() {
446        assert_eq!(49u64.sqrt_exact(), Some(7));
447        assert_eq!(50u64.sqrt_exact(), None);
448    }
449
450    #[test]
451    fn exact_roots_cbrt_exact() {
452        assert_eq!(125u64.cbrt_exact(), Some(5));
453        assert_eq!(126u64.cbrt_exact(), None);
454    }
455}