const_primes/generate.rs
1//! This module contains implementations of prime generation functions.
2
3use core::fmt;
4
5use crate::{sieve, sieve::sieve_segment, Underlying};
6
7/// Returns the `N` first prime numbers.
8///
9/// [`Primes`](crate::Primes) might be relevant for you if you intend to later use these prime numbers for related computations.
10///
11/// Uses a [segmented sieve of Eratosthenes](https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Segmented_sieve).
12///
13/// # Example
14///
15/// ```
16/// # use const_primes::primes;
17/// const PRIMES: [u32; 10] = primes();
18/// assert_eq!(PRIMES, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]);
19/// ```
20#[must_use = "the function only returns a new value"]
21pub const fn primes<const N: usize>() -> [Underlying; N] {
22 if N <= 1 {
23 return [2; N];
24 } else if N == 2 {
25 let mut primes = [0; N];
26 primes[0] = 2;
27 primes[1] = 3;
28 return primes;
29 }
30
31 // This is a segmented sieve that runs until it has found enough primes.
32
33 // This array is the output in the end
34 let mut primes = [0; N];
35 // This keeps track of how many primes we've found so far.
36 let mut prime_count = 0;
37
38 // Sieve the first primes below N
39 let mut sieve: [bool; N] = sieve();
40
41 // Count how many primes we found
42 // and store them in the final array
43 let mut number = 0;
44 while number < N {
45 if sieve[number] {
46 primes[prime_count] = number as Underlying;
47 prime_count += 1;
48 }
49
50 number += 1;
51 }
52
53 // For every segment of N numbers
54 let mut low = N - 1;
55 let mut high = 2 * N - 1;
56 'generate: while prime_count < N {
57 // reset the sieve for the segment
58 sieve = [true; N];
59 let mut i = 0;
60
61 // and repeat for each prime found so far:
62 while i < prime_count {
63 let prime = primes[i] as usize;
64
65 // Find the smallest composite in the current segment,
66 let mut composite = (low / prime) * prime;
67 if composite < low {
68 composite += prime;
69 }
70
71 // and sieve all numbers in the segment that are multiples of the prime.
72 while composite < high {
73 sieve[composite - low] = false;
74 composite += prime;
75 }
76
77 i += 1;
78 }
79
80 // Move the found primes into the final array
81 i = low;
82 while i < high {
83 if sieve[i - low] {
84 primes[prime_count] = i as Underlying;
85 prime_count += 1;
86 // and stop the generation of primes if we're done.
87 if prime_count >= N {
88 break 'generate;
89 }
90 }
91 i += 1;
92 }
93
94 // Update low and high for the next segment
95 low += N;
96 high += N;
97 }
98
99 primes
100}
101
102/// Returns the `N` largest primes less than `upper_limit`.
103///
104/// This function uses a segmented sieve of size `MEM` for computation,
105/// but only stores the `N` requested primes in the output array.
106///
107/// Set `MEM` such that `MEM*MEM >= upper_limit`.
108///
109/// If you want to compute primes that are larger than some limit, take a look at [`primes_geq`].
110///
111/// If you do not wish to compute the size requirement of the sieve manually, take a look at [`primes_segment!`](crate::primes_segment).
112///
113/// # Example
114///
115/// Basic usage:
116///
117/// ```
118/// # use const_primes::{primes_lt, GenerationError};
119/// // Sieving up to 100 means the sieve needs to be of size ceil(sqrt(100)) = 10.
120/// // However, we only save the 4 largest primes in the constant.
121/// const PRIMES: [u64;4] = match primes_lt::<4, 10>(100) {Ok(ps) => ps, Err(_) => panic!()};
122/// assert_eq!(PRIMES, [79, 83, 89, 97]);
123/// ```
124///
125/// Compute limited ranges of large primes. Functions provided by the crate can help you
126/// compute the needed sieve size:
127///
128/// ```
129/// # use const_primes::{primes_lt, GenerationError};
130/// use const_primes::isqrt;
131/// const N: usize = 3;
132/// const LIMIT: u64 = 5_000_000_030;
133/// const MEM: usize = isqrt(LIMIT) as usize + 1;
134/// const BIG_PRIMES: Result<[u64; N], GenerationError> = primes_lt::<N, MEM>(LIMIT);
135///
136/// assert_eq!(BIG_PRIMES, Ok([4_999_999_903, 4_999_999_937, 5_000_000_029]));
137/// ```
138///
139/// # Errors
140///
141/// If the number of primes requested, `N`, is larger than
142/// the number of primes that exists below the `upper_limit` this function
143/// returns an error:
144///
145/// ```
146/// # use const_primes::{primes_lt, GenerationError};
147/// const N: usize = 9;
148/// const PRIMES: Result<[u64; N], GenerationError> = primes_lt::<N, N>(10);
149/// assert_eq!(PRIMES, Err(GenerationError::OutOfPrimes));
150/// ```
151///
152/// It also returns an error if `upper_limit` is larger than `MEM`^2 or if `upper_limit` is smaller than or equal to 2:
153///
154/// ```
155/// # use const_primes::{primes_lt, GenerationError};
156/// const TOO_LARGE_LIMIT: Result<[u64; 3], GenerationError> = primes_lt::<3, 5>(26);
157/// const TOO_SMALL_LIMIT: Result<[u64; 1], GenerationError> = primes_lt::<1, 1>(1);
158/// assert_eq!(TOO_LARGE_LIMIT, Err(GenerationError::TooSmallSieveSize));
159/// assert_eq!(TOO_SMALL_LIMIT, Err(GenerationError::TooSmallLimit));
160/// ```
161///
162/// It is a compile error if `MEM` is smaller than `N`, or if `MEM`^2 does not fit in a `u64`:
163///
164/// ```compile_fail
165/// # use const_primes::{primes_lt, GenerationError};
166/// const TOO_SMALL_MEM: Result<[u64; 5], GenerationError> = primes_lt::<5, 2>(20);
167/// ```
168///
169/// ```compile_fail
170/// # use const_primes::{primes_lt, GenerationError};
171/// const TOO_BIG_MEM: Result<[u64; 10], GenerationError> = primes_lt::<10, 1_000_000_000_000>(100);
172/// ```
173#[must_use = "the function only returns a new value and does not modify its input"]
174pub const fn primes_lt<const N: usize, const MEM: usize>(
175 mut upper_limit: u64,
176) -> Result<[u64; N], GenerationError> {
177 const { assert!(MEM >= N, "`MEM` must be at least as large as `N`") }
178
179 let mem_sqr = const {
180 let mem64 = MEM as u64;
181 match mem64.checked_mul(mem64) {
182 Some(mem_sqr) => mem_sqr,
183 None => panic!("`MEM`^2 must fit in a u64"),
184 }
185 };
186
187 if upper_limit <= 2 {
188 return Err(GenerationError::TooSmallLimit);
189 }
190
191 if upper_limit > mem_sqr {
192 return Err(GenerationError::TooSmallSieveSize);
193 }
194
195 let mut primes: [u64; N] = [0; N];
196
197 if N == 0 {
198 return Ok(primes);
199 }
200
201 // This will be used to sieve all upper ranges.
202 let base_sieve: [bool; MEM] = sieve();
203
204 let mut total_primes_found: usize = 0;
205 'generate: while total_primes_found < N {
206 // This is the smallest prime we have found so far.
207 let mut smallest_found_prime = primes[N - 1 - total_primes_found];
208 // Sieve for primes in the segment.
209 let (offset, upper_sieve) = match sieve_segment(&base_sieve, upper_limit) {
210 Ok(res) => (0, res),
211 // The segment was larger than there are numbers left to sieve, just use the base sieve
212 Err(_) => ((MEM as u64 - upper_limit) as usize, base_sieve),
213 };
214
215 let mut i: usize = 0;
216 while i < MEM - offset {
217 // Iterate backwards through the upper sieve.
218 if upper_sieve[MEM - 1 - i - offset] {
219 smallest_found_prime = upper_limit - 1 - i as u64;
220 // Write every found prime to the primes array.
221 primes[N - 1 - total_primes_found] = smallest_found_prime;
222 total_primes_found += 1;
223 if total_primes_found >= N {
224 // If we have found enough primes we stop sieving.
225 break 'generate;
226 }
227 }
228 i += 1;
229 }
230 upper_limit = smallest_found_prime;
231 if upper_limit <= 2 && total_primes_found < N {
232 return Err(GenerationError::OutOfPrimes);
233 }
234 }
235
236 Ok(primes)
237}
238
239/// Generate arrays of large prime numbers without having to store all primes
240/// from 2 and up in the result, and thus potentially the binary.
241///
242/// Calls [`primes_geq`] or [`primes_lt`], and automatically computes the memory requirement of the sieve.
243///
244/// Compute `N` primes larger than or equal to some limit as `primes_segment!(N; >= LIMIT)`,
245/// and `N` primes less than some limit as `primes_segment!(N; < LIMIT)`.
246///
247/// Estimates the sieve size as `isqrt(upper_limit) + 1` for [`primes_lt`]
248/// and as `isqrt(lower_limit) + 1 + N` for [`primes_geq`].
249/// This may overestimate the memory requirement for `primes_geq`.
250///
251/// # Example
252///
253/// ```
254/// # use const_primes::{primes_segment, GenerationError};
255/// const N: usize = 3;
256/// const LIMIT: u64 = 5_000_000_031;
257///
258/// const PRIMES_GEQ: Result<[u64; N], GenerationError> = primes_segment!(N; >= LIMIT);
259/// const PRIMES_LT: Result<[u64; N], GenerationError> = primes_segment!(N; < LIMIT);
260///
261/// // Can also be used at runtime:
262/// let primes_geq = primes_segment!(N; >= LIMIT);
263///
264/// assert_eq!(PRIMES_GEQ, primes_geq);
265/// assert_eq!(PRIMES_GEQ, Ok([5000000039, 5000000059, 5000000063]));
266/// assert_eq!(PRIMES_LT, Ok([4999999903, 4999999937, 5000000029]));
267/// ```
268///
269/// # Errors
270///
271/// Has the same error behaviour as [`primes_geq`] and [`primes_lt`], with the exception
272/// that it sets `MEM` such that the sieve doesn't run out of memory.
273#[macro_export]
274macro_rules! primes_segment {
275 ($n:expr; < $lim:expr) => {
276 $crate::primes_lt::<
277 { $n },
278 {
279 let mem: u64 = { $lim };
280 $crate::isqrt(mem) as ::core::primitive::usize + 1
281 },
282 >({ $lim })
283 };
284 ($n:expr; >= $lim:expr) => {
285 $crate::primes_geq::<
286 { $n },
287 {
288 let mem: u64 = { $lim };
289 $crate::isqrt(mem) as ::core::primitive::usize + 1 + { $n }
290 },
291 >({ $lim })
292 };
293}
294
295/// Returns the `N` smallest primes greater than or equal to `lower_limit`.
296///
297/// This function uses a segmented sieve of size `MEM` for computation,
298/// but only stores the `N` requested primes in the output array.
299///
300/// Set `MEM` such that `MEM`^2 is larger than the largest prime you will encounter.
301///
302/// If you want to compute primes smaller than some limit, take a look at [`primes_lt`].
303///
304/// If you do not wish to compute the size requirement of the sieve manually, take a look at [`primes_segment!`](crate::primes_segment).
305///
306/// # Examples
307///
308/// Basic usage:
309///
310/// ```
311/// use const_primes::primes_geq;
312/// // Compute 5 primes larger than 40. The largest will be 59, so `MEM` needs to be at least 8.
313/// const PRIMES: [u64; 5] = match primes_geq::<5, 8>(40) {Ok(ps) => ps, Err(_) => panic!()};
314/// assert_eq!(PRIMES, [41, 43, 47, 53, 59]);
315/// ```
316///
317/// Compute limited ranges of large primes. Functions provided by the crate can help you
318/// compute the needed sieve size:
319///
320/// ```
321/// # use const_primes::{primes_geq, GenerationError};
322/// use const_primes::isqrt;
323/// const N: usize = 3;
324/// const LIMIT: u64 = 5_000_000_030;
325/// const MEM: usize = isqrt(LIMIT) as usize + 1 + N;
326/// const PRIMES_GEQ: Result<[u64; N], GenerationError> = primes_geq::<N, MEM>(LIMIT);
327/// assert_eq!(PRIMES_GEQ, Ok([5_000_000_039, 5_000_000_059, 5_000_000_063]));
328/// # Ok::<(), GenerationError>(())
329/// ```
330///
331/// # Errors
332///
333/// Only primes smaller than `MEM^2` can be generated, so if the sieve
334/// encounters a number larger than that it results in an error:
335///
336/// ```
337/// # use const_primes::{primes_geq, GenerationError};
338/// const PRIMES: Result<[u64; 3], GenerationError> = primes_geq::<3, 3>(5);
339/// // The sieve is unable to determine the prime status of 9,
340/// // since that is the same or larger than `MEM`^2.
341/// assert_eq!(PRIMES, Err(GenerationError::SieveOverrun(9)));
342/// ```
343///
344/// Also returns an error if `lower_limit` is larger than or equal to `MEM^2`:
345///
346/// ```
347/// # use const_primes::{primes_geq, GenerationError};
348/// const PRIMES: Result<[u64; 5], GenerationError> = primes_geq::<5, 5>(26);
349/// assert_eq!(PRIMES, Err(GenerationError::TooSmallSieveSize));
350/// ```
351///
352/// It is a compile error if `MEM` is smaller than `N`, or if `MEM`^2 does not fit in a `u64`:
353///
354/// ```compile_fail
355/// # use const_primes::{primes_geq, GenerationError};
356/// const TOO_SMALL_MEM: Result<[u64; 5], GenerationError> = primes_geq::<5, 2>(20);
357/// ```
358///
359/// ```compile_fail
360/// # use const_primes::{primes_geq, GenerationError};
361/// const TOO_BIG_MEM: Result<[u64; 10], GenerationError> = primes_geq::<10, 1_000_000_000_000>(100);
362/// ```
363#[must_use = "the function only returns a new value and does not modify its input"]
364pub const fn primes_geq<const N: usize, const MEM: usize>(
365 lower_limit: u64,
366) -> Result<[u64; N], GenerationError> {
367 const { assert!(MEM >= N, "`MEM` must be at least as large as `N`") }
368
369 let (mem64, mem_sqr) = const {
370 let mem64 = MEM as u64;
371 match mem64.checked_mul(mem64) {
372 Some(mem_sqr) => (mem64, mem_sqr),
373 None => panic!("`MEM`^2 must fit in a `u64`"),
374 }
375 };
376
377 if N == 0 {
378 return Ok([0; N]);
379 }
380
381 // If `lower_limit` is 2 or less, this is the same as calling `primes`,
382 // so we just do that and convert the result to `u64`.
383 if lower_limit <= 2 {
384 let ans32: [u32; N] = primes();
385 let mut ans64 = [0; N];
386 let mut i = 0;
387 while i < N {
388 ans64[i] = ans32[i] as u64;
389 i += 1;
390 }
391 return Ok(ans64);
392 }
393
394 if lower_limit >= mem_sqr {
395 return Err(GenerationError::TooSmallSieveSize);
396 }
397
398 let mut primes = [0; N];
399 let mut total_found_primes = 0;
400 let mut largest_found_prime = 0;
401 let base_sieve: [bool; MEM] = sieve();
402 let mut sieve_limit = lower_limit;
403 'generate: while total_found_primes < N {
404 let Ok(upper_sieve) = sieve_segment(&base_sieve, sieve_limit + mem64) else {
405 panic!("can not happen since we set upper limit to mem + nonzero stuff")
406 };
407
408 let mut i = 0;
409 while i < MEM {
410 if upper_sieve[i] {
411 largest_found_prime = sieve_limit + i as u64;
412
413 // We can not know whether this is actually a prime since
414 // the base sieve contains no information
415 // about numbers larger than or equal to `MEM`^2.
416 if largest_found_prime >= mem_sqr {
417 return Err(GenerationError::SieveOverrun(largest_found_prime));
418 }
419
420 if largest_found_prime >= lower_limit {
421 primes[total_found_primes] = largest_found_prime;
422 total_found_primes += 1;
423 if total_found_primes >= N {
424 // We've found enough primes.
425 break 'generate;
426 }
427 }
428 }
429 i += 1;
430 }
431 sieve_limit = largest_found_prime + 1;
432 }
433
434 Ok(primes)
435}
436
437/// The error returned by [`primes_lt`] and [`primes_geq`] if the input
438/// is invalid or does not work to produce the requested primes.
439#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
440#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
441#[cfg_attr(
442 feature = "rkyv",
443 derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)
444)]
445pub enum GenerationError {
446 /// The limit was larger than or equal to `MEM^2`.
447 TooSmallSieveSize,
448 /// The limit was smaller than or equal to 2.
449 TooSmallLimit,
450 /// Encountered a number larger than or equal to `MEM`^2.
451 SieveOverrun(u64),
452 /// Ran out of primes.
453 OutOfPrimes,
454}
455
456impl fmt::Display for GenerationError {
457 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
458 match self {
459 Self::TooSmallSieveSize => write!(
460 f,
461 "the limit was larger than `MEM`^2"
462 ),
463 Self::TooSmallLimit => write!(
464 f,
465 "the limit was smaller than or equal to 2"
466 ),
467 Self::SieveOverrun(number) => write!(
468 f,
469 "encountered the number {number} which would have needed `MEM` to be at least {} to sieve", crate::integer_math::isqrt(*number) + 1
470 ),
471 Self::OutOfPrimes => write!(f, "ran out of primes before the array was filled"),
472 }
473 }
474}
475
476impl core::error::Error for GenerationError {}
477
478#[cfg(test)]
479mod test {
480 use crate::is_prime;
481
482 use super::*;
483
484 #[test]
485 fn sanity_check_primes_geq() {
486 {
487 const P: Result<[u64; 5], GenerationError> = primes_geq::<5, 5>(10);
488 assert_eq!(P, Ok([11, 13, 17, 19, 23]));
489 }
490 {
491 const P: Result<[u64; 5], GenerationError> = primes_geq::<5, 5>(0);
492 assert_eq!(P, Ok([2, 3, 5, 7, 11]));
493 }
494 {
495 const P: Result<[u64; 1], GenerationError> = primes_geq::<1, 1>(0);
496 assert_eq!(P, Ok([2]));
497 }
498 for &prime in primes_geq::<2_000, 2_008>(3_998_000).unwrap().as_slice() {
499 assert!(is_prime(prime));
500 }
501 assert_eq!(primes_geq::<0, 0>(10), Ok([]));
502 assert_eq!(primes_geq::<3, 3>(2), Ok([2, 3, 5]));
503 assert_eq!(
504 primes_geq::<3, 3>(10),
505 Err(GenerationError::TooSmallSieveSize)
506 );
507 assert_eq!(primes_geq::<2, 2>(3), Err(GenerationError::SieveOverrun(4)));
508 }
509
510 #[test]
511 fn sanity_check_primes_lt() {
512 {
513 const P: Result<[u64; 5], GenerationError> = primes_lt::<5, 5>(20);
514 assert_eq!(P, Ok([7, 11, 13, 17, 19]));
515 }
516 {
517 const P: Result<[u64; 5], GenerationError> = primes_lt::<5, 5>(12);
518 assert_eq!(P, Ok([2, 3, 5, 7, 11]));
519 }
520 {
521 const P: Result<[u64; 1], GenerationError> = primes_lt::<1, 2>(3);
522 assert_eq!(P, Ok([2]));
523 }
524 assert_eq!(primes_lt::<2, 2>(2), Err(GenerationError::TooSmallLimit));
525 assert_eq!(
526 primes_lt::<2, 2>(5),
527 Err(GenerationError::TooSmallSieveSize)
528 );
529 assert_eq!(primes_lt::<0, 2>(3), Ok([]));
530 assert_eq!(primes_lt::<3, 5>(4), Err(GenerationError::OutOfPrimes));
531 }
532
533 #[test]
534 fn check_primes_segment() {
535 const P_GEQ: Result<[u64; 10], GenerationError> = primes_segment!(10; >= 1000);
536 const P_LT: Result<[u64; 10], GenerationError> = primes_segment!(10; < 1000);
537
538 assert_eq!(
539 P_GEQ,
540 Ok([1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061])
541 );
542 assert_eq!(P_LT, Ok([937, 941, 947, 953, 967, 971, 977, 983, 991, 997]));
543 }
544}