prime_iter/lib.rs
1#![cfg_attr(not(feature = "std"), no_std)]
2//! A somewhat optimized incremental-sieve based prime generator.
3//!
4//! ## Examples
5//! The interface is given in the form of an iterator, so usage is very simple and idiomatic:
6//! ```
7//! let fifty_second_prime = prime_iter::primes::<i32>().nth(51).unwrap();
8//!
9//! assert_eq!(fifty_second_prime, 239);
10//! ```
11//! ```
12//! let prime_sum: i32 = prime_iter::primes::<i32>().take(100).sum();
13//!
14//! assert_eq!(prime_sum, 24133);
15//! ```
16//! ```
17//! let two_digit_primes: Vec<i32> = prime_iter::primes::<i32>().skip_while(|&x| x < 10).take_while(|&x| x < 100).collect();
18//!
19//! assert_eq!(two_digit_primes, [11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]);
20//! ```
21//! And of course `for` loops work too:
22//! ```no_run
23//! for prime in prime_iter::primes::<i32>() {
24//! if prime % 10 == 1 {
25//! println!("{prime}");
26//! }
27//! }
28//! ```
29//!
30//! ## `no_std` Support
31//! `prime-iter` supports no-std environments, however it does use allocations. Disable the `std` feature,
32//! which is enabled by default, for a no_std environment.
33
34#[cfg(not(feature = "std"))]
35extern crate alloc;
36#[cfg(not(feature = "std"))]
37use alloc::vec::Vec;
38
39/// Returns an iterator that provides all of the primes that fit into `T`, in increasing order.
40///
41/// `T` may only be a primitive integer type (`iXX`, `uXX` and `isize`, `usize`)
42pub fn primes<T>() -> Primes<T>
43where
44 Primes<T>: Default + Iterator<Item = T>,
45{
46 Primes::default()
47}
48
49/// The iterator type returned by [primes].
50#[derive(Clone, Debug)]
51pub struct Primes<T> {
52 primes: Vec<T>,
53 multiples: Vec<T>,
54 current: T,
55 done: bool,
56}
57
58macro_rules! impl_primes {
59 ($ty:path) => {
60 impl Default for Primes<$ty> {
61 fn default() -> Self {
62 Self {
63 primes: Default::default(),
64 multiples: Default::default(),
65 current: 2,
66 done: false,
67 }
68 }
69 }
70
71 impl Iterator for Primes<$ty> {
72 type Item = $ty;
73
74 fn next(&mut self) -> Option<Self::Item> {
75 if self.done {
76 return None;
77 }
78 if self.current == 2 {
79 self.current += 1;
80 return Some(2);
81 }
82
83 loop {
84 let mut is_prime = true;
85
86 'primeloop: for (i, &prime) in self.primes.iter().enumerate() {
87 if self.multiples[i] == 0 {
88 continue;
89 }
90 while self.multiples[i] < self.current {
91 match self.multiples[i].checked_add(prime) {
92 Some(x) => self.multiples[i] = x,
93 None => {
94 self.multiples[i] = 0;
95 continue 'primeloop;
96 }
97 }
98 }
99
100 if self.multiples[i] == self.current {
101 is_prime = false;
102 break;
103 }
104 }
105
106 let current = self.current;
107
108 if is_prime {
109 if let Some(res) = self.current.checked_mul(self.current) {
110 self.primes.push(self.current);
111 self.multiples.push(res);
112 }
113 }
114
115 match self.current.checked_add(2) {
116 Some(num) => self.current = num,
117 None => {
118 self.done = true;
119 break None;
120 }
121 }
122
123 if is_prime {
124 break Some(current);
125 }
126 }
127 }
128 }
129 };
130
131 ($($ty:path),+) => {
132 $(
133 impl_primes!($ty);
134 )+
135 };
136}
137
138impl_primes!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);