1use crate::factor::{pollard_rho, trial_division};
14use crate::nt_funcs::{
15 factorize128, is_prime64, next_prime, nth_prime_bounds, nth_prime_est, prev_prime,
16};
17use crate::primality::{PrimalityBase, PrimalityRefBase};
18use crate::tables::{SMALL_PRIMES, SMALL_PRIMES_NEXT};
19use crate::traits::{
20 FactorizationConfig, Primality, PrimalityTestConfig, PrimalityUtils, PrimeBuffer,
21};
22use bitvec::{bitvec, prelude::Msb0};
23use lru::LruCache;
24use num_integer::Roots;
25use rand::random;
26use std::collections::BTreeMap;
27use std::num::NonZeroUsize;
28
29pub trait PrimeBufferExt: for<'a> PrimeBuffer<'a> {
31 fn is_prime<T: PrimalityBase>(
40 &self,
41 target: &T,
42 config: Option<PrimalityTestConfig>,
43 ) -> Primality
44 where
45 for<'r> &'r T: PrimalityRefBase<T>,
46 {
47 if target.is_even() {
49 return if target == &T::from_u8(2u8).unwrap() {
50 Primality::Yes
51 } else {
52 Primality::No
53 };
54 }
55
56 if let Some(x) = target.to_u64() {
58 return match is_prime64(x) {
59 true => Primality::Yes,
60 false => Primality::No,
61 };
62 }
63
64 let config = config.unwrap_or(PrimalityTestConfig::default());
65 let mut probability = 1.;
66
67 let mut witness_list: Vec<u64> = Vec::new();
69 if config.sprp_trials > 0 {
70 witness_list.extend(self.iter().take(config.sprp_trials));
71 probability *= 1. - 0.25f32.powi(config.sprp_trials as i32);
72 }
73 if config.sprp_random_trials > 0 {
74 for _ in 0..config.sprp_random_trials {
75 let mut w: u64 = rand::random();
77 while witness_list.iter().find(|&x| x == &w).is_some() {
78 w = rand::random();
79 }
80 witness_list.push(w);
81 }
82 probability *= 1. - 0.25f32.powi(config.sprp_random_trials as i32);
83 }
84 if !witness_list
85 .into_iter()
86 .all(|x| target.is_sprp(T::from_u64(x).unwrap()))
87 {
88 return Primality::No;
89 }
90
91 if config.slprp_test {
93 probability *= 1. - 4f32 / 15f32;
94 if !target.is_slprp(None, None) {
95 return Primality::No;
96 }
97 }
98 if config.eslprp_test {
99 probability *= 1. - 4f32 / 15f32;
100 if !target.is_eslprp(None) {
101 return Primality::No;
102 }
103 }
104
105 Primality::Probable(probability)
106 }
107
108 fn factors<T: PrimalityBase>(
118 &self,
119 target: T,
120 config: Option<FactorizationConfig>,
121 ) -> (BTreeMap<T, usize>, Option<Vec<T>>)
122 where
123 for<'r> &'r T: PrimalityRefBase<T>,
124 {
125 if let Some(x) = target.to_u128() {
127 let factors = factorize128(x)
128 .into_iter()
129 .map(|(k, v)| (T::from_u128(k).unwrap(), v))
130 .collect();
131 return (factors, None);
132 }
133 let config = config.unwrap_or(FactorizationConfig::default());
134
135 let (result, factored) = trial_division(self.iter().cloned(), target, config.td_limit);
137 let mut result: BTreeMap<T, usize> = result
138 .into_iter()
139 .map(|(k, v)| (T::from_u64(k).unwrap(), v))
140 .collect();
141
142 let mut failed = Vec::new();
146 let mut config = config;
147 config.td_limit = Some(0); match factored {
149 Ok(res) => {
150 if !res.is_one() {
151 result.insert(res, 1);
152 }
153 }
154 Err(res) => {
155 let mut todo = vec![res];
156 while let Some(target) = todo.pop() {
157 if self
158 .is_prime(&target, Some(config.primality_config))
159 .probably()
160 {
161 *result.entry(target).or_insert(0) += 1;
162 } else {
163 if let Some(divisor) = self.divisor(&target, &mut config) {
164 todo.push(divisor.clone());
165 todo.push(target / divisor);
166 } else {
167 failed.push(target);
168 }
169 }
170 }
171 }
172 };
173
174 if failed.is_empty() {
175 (result, None)
176 } else {
177 (result, Some(failed))
178 }
179 }
180
181 fn factorize<T: PrimalityBase>(&self, target: T) -> BTreeMap<T, usize>
186 where
187 for<'r> &'r T: PrimalityRefBase<T>,
188 {
189 loop {
191 let (result, remainder) = self.factors(target.clone(), None);
192 if remainder.is_none() {
193 break result;
194 }
195 }
196 }
197
198 fn divisor<T: PrimalityBase>(&self, target: &T, config: &mut FactorizationConfig) -> Option<T>
203 where
204 for<'r> &'r T: PrimalityRefBase<T>,
205 {
206 if matches!(config.td_limit, Some(0)) {
207 let tsqrt: T = Roots::sqrt(target) + T::one();
209 let limit = if let Some(l) = config.td_limit {
210 tsqrt.clone().min(T::from_u64(l).unwrap())
211 } else {
212 tsqrt.clone()
213 };
214
215 for p in self.iter().map(|p| T::from_u64(*p).unwrap()) {
216 if &p > &tsqrt {
217 return None; }
219 if &p > &limit {
220 break;
221 }
222 if target.is_multiple_of(&p) {
223 return Some(p);
224 }
225 }
226 }
227
228 let below64 = target.to_u64().is_some();
230 while config.rho_trials > 0 {
231 let (start, offset) = if below64 {
232 (
233 T::from_u8(random::<u8>()).unwrap() % target,
234 T::from_u8(random::<u8>()).unwrap() % target,
235 )
236 } else {
237 (
238 T::from_u64(random::<u64>()).unwrap() % target,
239 T::from_u64(random::<u64>()).unwrap() % target,
240 )
241 };
242 config.rho_trials -= 1;
243 if let (Some(p), _) = pollard_rho(target, start, offset, 1048576) {
246 return Some(p);
247 }
248 }
249
250 None
251 }
252}
253
254impl<T> PrimeBufferExt for T where for<'a> T: PrimeBuffer<'a> {}
255
256pub struct NaiveBuffer {
258 list: Vec<u64>, next: u64, }
261
262impl NaiveBuffer {
263 #[inline]
264 pub fn new() -> Self {
265 let list = SMALL_PRIMES.iter().map(|&p| p as u64).collect();
266 NaiveBuffer {
267 list,
268 next: SMALL_PRIMES_NEXT,
269 }
270 }
271}
272
273impl<'a> PrimeBuffer<'a> for NaiveBuffer {
274 type PrimeIter = std::slice::Iter<'a, u64>;
275
276 fn contains(&self, num: u64) -> bool {
277 self.list.binary_search(&num).is_ok()
278 }
279
280 fn clear(&mut self) {
281 self.list.truncate(16);
282 self.list.shrink_to_fit();
283 self.next = 55; }
285
286 fn iter(&'a self) -> Self::PrimeIter {
287 self.list.iter()
288 }
289
290 fn bound(&self) -> u64 {
291 *self.list.last().unwrap()
292 }
293
294 fn reserve(&mut self, limit: u64) {
295 let sieve_limit = (limit | 1) + 2; let current = self.next; debug_assert!(current % 2 == 1);
298 if sieve_limit < current {
299 return;
300 }
301
302 let mut sieve = bitvec![usize, Msb0; 0; ((sieve_limit - current) / 2) as usize];
304 for p in self.list.iter().skip(1) {
305 let start = if p * p < current {
307 p * ((current / p) | 1) } else {
309 p * p
310 };
311 for multi in (start..sieve_limit).step_by(2 * (*p as usize)) {
312 if multi >= current {
313 sieve.set(((multi - current) / 2) as usize, true);
314 }
315 }
316 }
317
318 for p in (current..Roots::sqrt(&sieve_limit) + 1).step_by(2) {
320 for multi in (p * p..sieve_limit).step_by(2 * (p as usize)) {
321 if multi >= current {
322 sieve.set(((multi - current) / 2) as usize, true);
323 }
324 }
325 }
326
327 self.list
329 .extend(sieve.iter_zeros().map(|x| (x as u64) * 2 + current));
330 self.next = sieve_limit;
331 }
332}
333
334impl NaiveBuffer {
335 pub fn primorial<T: PrimalityBase + std::iter::Product>(&mut self, n: usize) -> T {
340 self.nprimes(n).map(|&p| T::from_u64(p).unwrap()).product()
341 }
342
343 pub fn primes(&mut self, limit: u64) -> std::iter::Take<<Self as PrimeBuffer>::PrimeIter> {
349 self.reserve(limit);
350 let position = match self.list.binary_search(&limit) {
351 Ok(p) => p + 1,
352 Err(p) => p,
353 }; return self.list.iter().take(position);
355 }
356
357 pub fn into_primes(mut self, limit: u64) -> std::vec::IntoIter<u64> {
359 self.reserve(limit);
360 let position = match self.list.binary_search(&limit) {
361 Ok(p) => p + 1,
362 Err(p) => p,
363 }; self.list.truncate(position);
365 return self.list.into_iter();
366 }
367
368 pub fn nprimes(&mut self, count: usize) -> std::iter::Take<<Self as PrimeBuffer>::PrimeIter> {
370 let (_, bound) = nth_prime_bounds(&(count as u64))
371 .expect("Estimated size of the largest prime will be larger than u64 limit");
372 self.reserve(bound);
373 debug_assert!(self.list.len() >= count);
374 self.list.iter().take(count)
375 }
376
377 pub fn into_nprimes(mut self, count: usize) -> std::vec::IntoIter<u64> {
379 let (_, bound) = nth_prime_bounds(&(count as u64))
380 .expect("Estimated size of the largest prime will be larger than u64 limit");
381 self.reserve(bound);
382 debug_assert!(self.list.len() >= count);
383 self.list.truncate(count);
384 return self.list.into_iter();
385 }
386
387 pub fn nth_prime(&mut self, n: u64) -> u64 {
392 if n < self.list.len() as u64 {
393 return self.list[n as usize - 1];
394 }
395
396 const THRESHOLD_NTH_PRIME_SIEVE: u64 = 4096;
398 if n <= THRESHOLD_NTH_PRIME_SIEVE {
399 return *self.nprimes(n as usize).last().unwrap();
400 }
401
402 let mut x = prev_prime(&nth_prime_est(&n).unwrap(), None).unwrap();
404 let mut pi = self.prime_pi(x);
405
406 while pi > n {
407 x = prev_prime(&x, None).unwrap();
408 pi -= 1;
409 }
410 while pi < n {
411 x = next_prime(&x, None).unwrap();
412 pi += 1;
413 }
414 x
415 }
416
417 pub fn prime_phi(&mut self, x: u64, a: usize, cache: &mut LruCache<(u64, usize), u64>) -> u64 {
419 if a == 1 {
420 return (x + 1) / 2;
421 }
422 if let Some(v) = cache.get(&(x, a)) {
423 return *v;
424 }
425 let t1 = self.prime_phi(x, a - 1, cache);
426 let pa = self.nth_prime(a as u64);
427 let t2 = self.prime_phi(x / pa, a - 1, cache);
428 let t = t1 - t2;
429 cache.put((x, a), t);
430 t
431 }
432
433 pub fn prime_pi(&mut self, limit: u64) -> u64 {
437 const THRESHOLD_PRIME_PI_SIEVE: u64 = 38873; if &limit <= self.list.last().unwrap() || limit <= THRESHOLD_PRIME_PI_SIEVE {
440 self.reserve(limit);
441
442 return match self.list.binary_search(&limit) {
443 Ok(pos) => (pos + 1) as u64,
444 Err(pos) => pos as u64,
445 };
446 }
447
448 let b = limit.sqrt();
450 let a = b.sqrt();
451 let c = limit.cbrt();
452 self.reserve(b);
453
454 let a = self.prime_pi(a);
455 let b = self.prime_pi(b);
456 let c = self.prime_pi(c);
457
458 let cache_cap = NonZeroUsize::new(a as usize).unwrap(); let mut phi_cache = LruCache::new(cache_cap);
460 let mut sum =
461 self.prime_phi(limit, a as usize, &mut phi_cache) + (b + a - 2) * (b - a + 1) / 2;
462 for i in a + 1..b + 1 {
463 let w = limit / self.nth_prime(i);
464 sum -= self.prime_pi(w);
465 if i <= c {
466 let l = self.prime_pi(w.sqrt());
467 sum += (l * (l - 1) - i * (i - 3)) / 2 - 1;
468 for j in i..(l + 1) {
469 let pj = self.nth_prime(j);
470 sum -= self.prime_pi(w / pj);
471 }
472 }
473 }
474 return sum;
475 }
476}
477
478#[cfg(test)]
479mod tests {
480 use super::*;
481 use crate::mint::SmallMint;
482 #[cfg(feature = "num-bigint")]
483 use core::str::FromStr;
484 #[cfg(feature = "num-bigint")]
485 use num_bigint::BigUint;
486 use rand::random;
487
488 #[test]
489 fn prime_generation_test() {
490 const PRIME50: [u64; 15] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47];
491 const PRIME300: [u64; 62] = [
492 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83,
493 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179,
494 181, 191, 193, 197, 199, 211, 223, 227, 229, 233, 239, 241, 251, 257, 263, 269, 271,
495 277, 281, 283, 293,
496 ];
497
498 let mut pb = NaiveBuffer::new();
499 assert_eq!(pb.primes(50).cloned().collect::<Vec<_>>(), PRIME50);
500 assert_eq!(pb.primes(300).cloned().collect::<Vec<_>>(), PRIME300);
501
502 pb.clear();
504 assert_eq!(pb.primes(293).cloned().collect::<Vec<_>>(), PRIME300);
505 pb = NaiveBuffer::new();
506 assert_eq!(*pb.primes(257).last().unwrap(), 257); pb = NaiveBuffer::new();
508 assert_eq!(*pb.primes(8167).last().unwrap(), 8167); }
510
511 #[test]
512 fn nth_prime_test() {
513 let mut pb = NaiveBuffer::new();
514 assert_eq!(pb.nth_prime(10000), 104729);
515 assert_eq!(pb.nth_prime(20000), 224737);
516 assert_eq!(pb.nth_prime(10000), 104729); assert_eq!(pb.nth_prime(10u64.pow(4)), 104729);
520 assert_eq!(pb.nth_prime(10u64.pow(5)), 1299709);
521 assert_eq!(pb.nth_prime(10u64.pow(6)), 15485863);
522 assert_eq!(pb.nth_prime(10u64.pow(7)), 179424673);
523 }
524
525 #[test]
526 fn prime_pi_test() {
527 let mut pb = NaiveBuffer::new();
528 assert_eq!(pb.prime_pi(8161), 1024); assert_eq!(pb.prime_pi(10000), 1229);
530 assert_eq!(pb.prime_pi(20000), 2262);
531 assert_eq!(pb.prime_pi(38873), 4096);
532
533 pb.clear(); assert_eq!(pb.prime_pi(8161), 1024);
535 assert_eq!(pb.prime_pi(10000), 1229);
536 assert_eq!(pb.prime_pi(20000), 2262);
537 assert_eq!(pb.prime_pi(10000), 1229); assert_eq!(pb.prime_pi(10u64.pow(5)), 9592);
541 assert_eq!(pb.prime_pi(10u64.pow(6)), 78498);
542 assert_eq!(pb.prime_pi(10u64.pow(7)), 664579);
543 assert_eq!(pb.prime_pi(10u64.pow(8)), 5761455);
544 }
545
546 #[test]
547 fn is_prime_test() {
548 let pb = NaiveBuffer::new();
550
551 assert_eq!(pb.is_prime(&(2u32.pow(19) - 1), None), Primality::Yes);
553 assert_eq!(pb.is_prime(&(2u32.pow(23) - 1), None), Primality::No);
554 assert!(matches!(
555 pb.is_prime(&(2u128.pow(89) - 1), None),
556 Primality::Probable(_)
557 ));
558 assert!(matches!(
559 pb.is_prime(&SmallMint::from(2u128.pow(89) - 1), None),
560 Primality::Probable(_)
561 ));
562
563 for _ in 0..100 {
565 let target = random::<u64>();
566 assert_eq!(
567 !is_prime64(target),
568 matches!(
569 pb.is_prime(&target, Some(PrimalityTestConfig::bpsw())),
570 Primality::No
571 )
572 );
573 }
574
575 const P: u128 = 18699199384836356663; assert!(matches!(pb.is_prime(&P, None), Primality::Probable(_)));
578 assert!(matches!(
579 pb.is_prime(&P, Some(PrimalityTestConfig::bpsw())),
580 Primality::Probable(_)
581 ));
582 assert!(matches!(
583 pb.is_prime(&SmallMint::from(P), None),
584 Primality::Probable(_)
585 ));
586 assert!(matches!(
587 pb.is_prime(&SmallMint::from(P), Some(PrimalityTestConfig::bpsw())),
588 Primality::Probable(_)
589 ));
590
591 const P2: u128 = 2019922777445599503530083;
592 assert!(matches!(pb.is_prime(&P2, None), Primality::Probable(_)));
593 assert!(matches!(
594 pb.is_prime(&P2, Some(PrimalityTestConfig::bpsw())),
595 Primality::Probable(_)
596 ));
597 assert!(matches!(
598 pb.is_prime(&SmallMint::from(P2), None),
599 Primality::Probable(_)
600 ));
601 assert!(matches!(
602 pb.is_prime(&SmallMint::from(P2), Some(PrimalityTestConfig::bpsw())),
603 Primality::Probable(_)
604 ));
605
606 #[cfg(feature = "num-bigint")]
607 {
608 let large_primes = [
609 "98920366548084643601728869055592650835572950932266967461790948584315647051443",
610 "94560208308847015747498523884063394671606671904944666360068158221458669711639",
611
612 "449417999055441493994709297093108513015373787049558499205492347871729927573118262811508386655998299074566974373711472560655026288668094291699357843464363003144674940345912431129144354948751003607115263071543163",
614 "230975859993204150666423538988557839555560243929065415434980904258310530753006723857139742334640122533598517597674807096648905501653461687601339782814316124971547968912893214002992086353183070342498989426570593",
615 "5521712099665906221540423207019333379125265462121169655563495403888449493493629943498064604536961775110765377745550377067893607246020694972959780839151452457728855382113555867743022746090187341871655890805971735385789993",
616 "203956878356401977405765866929034577280193993314348263094772646453283062722701277632936616063144088173312372882677123879538709400158306567338328279154499698366071906766440037074217117805690872792848149112022286332144876183376326512083574821647933992961249917319836219304274280243803104015000563790123",
617 "3618502788666131106986593281521497120414687020801267626233049500247285301239", "57896044618658097711785492504343953926634992332820282019728792003956564819949", "9850501549098619803069760025035903451269934817616361666987073351061430442874302652853566563721228910201656997576599", "42307582002575910332922579714097346549017899709713998034217522897561970639123926132812109468141778230245837569601494931472367", "6864797660130609714981900799081393217269435300143305409394463459185543183397656052122559640661454554977296311391480858037121987999716643812574028291115057151", "169511182982703321453314585423962898651587669459838234386506572286328885534468792292646838949809616446341407457141008401355628947670484184607678853094537849610289912805960069455687743151708433319901176932959509872662610091644590437761688516626993416011399330087939042347256922771590903190536793274742859624657"
626 ];
627 for pstr in large_primes {
628 assert!(
629 pb.is_prime(&BigUint::from_str(pstr).unwrap(), None)
630 .probably(),
631 "false negative prime {}",
632 pstr
633 )
634 }
635
636 let large_composites = [
637 "21284175091214687912771199898307297748211672914763848041968395774954376176754",
638 "6084766654921918907427900243509372380954290099172559290432744450051395395951",
639 "84594350493221918389213352992032324280367711247940675652888030554255915464401",
640 "82793403787388584738507275144194252681",
641
642 "1195068768795265792518361315725116351898245581", "8038374574536394912570796143419421081388376882875581458374889175222974273765333652186502336163960045457915042023603208766569966760987284043965408232928738791850869166857328267761771029389697739470167082304286871099974399765441448453411558724506334092790222752962294149842306881685404326457534018329786111298960644845216191652872597534901",
647 ];
648 for cstr in large_composites {
649 assert!(
650 !pb.is_prime(
651 &BigUint::from_str(cstr).unwrap(),
652 Some(PrimalityTestConfig::bpsw())
653 )
654 .probably(),
655 "false positive prime {}",
656 cstr
657 )
658 }
659 }
660 }
661
662 #[test]
663 fn pb_factors_test() {
664 let pb = NaiveBuffer::new();
665
666 #[cfg(feature = "num-bigint")]
667 {
668 let m131 = BigUint::from(2u8).pow(131) - 1u8; let (fac, r) = pb.factors(m131, None);
670 assert!(fac.len() == 2 && r.is_none());
671 }
672 }
673}