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