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);