I wrote a prime number sieve in Ruby last year as part of a coding challenge, but lately I wanted to port it to Rust. I finally got around to it, but it's still intensely Rubinic.
Before I take the obvious next step and parallelize it, I want to make sure my code is as Rust-idiomatic as possible, so that I don't develop bad Rust habits.
Docs and tests can be run with Cargo from the source repository.
Code
It consists of three parts: a SquareMultiple iterator that computes the sequence \$i^2, i^2+i, i^2 + 2i...\$, a sieve of Eratosthenes, and then a segmented sieve that uses multiples from the Eratosthenes sieve to sieve the rest of the numbers up to the first function argument.
SquareMultiple
The iterator code came from an answer to this SO question.
/// The `prime-suspects` crate provides functions that sieve primes.
extern crate itertools;
use itertools::Itertools;
/// A struct that is used to generate the sequence j = i^2, i^2+i, i^2+2i,
/// i^2+3i, ..., used in the sieve of Eratosthenes.
struct SquareMultiple {
/// The current iterator value. Starts at `arg * arg`, where `arg` is the
/// single argument given to the constructor.
curr: usize,
/// The increment is just `arg`. It is added each time the iterator's .next()
/// method is called.
inc: usize
}
/// An `Iterator` implementation of `SquareMultiple`.
impl Iterator for SquareMultiple {
type Item = usize;
/// The return type is `Option<T>`:
/// * When the `Iterator` is finished, `None` is returned.
/// * Otherwise, the next value is wrapped in `Some` and returned.
///
/// Since there's no endpoint to this sequence, this specific `Iterator` will
/// never return `None`. `Some` is always returned.
///
/// The iterator uses only `.curr` and `.inc`, because all there is to do is
/// add another of the original value (initialized in the implementation).
fn next(&mut self) -> Option<usize> {
let val = self.curr;
self.curr += self.inc;
Some(val)
}
}
/// The SquareMultiple implementation itself, using the struct and Iterator
/// traits defined above.
impl SquareMultiple {
fn new(term: usize) -> Self {
SquareMultiple { curr: term * term, inc: term }
}
}
Eratosthenes' Sieve
This code was a near-straight implementation of the Ruby version, which follows the pseudocode in the Wikipedia article.
/// An implementation of the sieve of Eratosthenes, as described in [the
/// Wikipedia
/// article](https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Algorithm_and_variants).
///
/// # Examples
/// ```
/// assert_eq!(vec![2,3,5,7], prime_suspects::eratosthenes_sieve(10));
/// ```
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::eratosthenes_sieve(65535)
/// .iter().rev().take(3).map(|&num| num)
/// .collect::<Vec<usize>>());
/// ```
pub fn eratosthenes_sieve(max_val: usize) -> Vec<usize> {
// Algorithm notes: The sieve works like so (Wikipedia, pseudocode inlined):
// Input: an integer n > 1
// Let A be an array of Boolean values, indexed by integers 2 to n,
// initially all set to true.
let mut bool_vec = vec![true; max_val];
// for i = 2, 3, 4, ..., not exceeding √n:
let mut top_sieve = max_val as f64;
// We have to add 1 because the sqrt coerced to an int is √floor(n)
top_sieve = top_sieve.sqrt() + 1.0;
for sieve_term in 2..(top_sieve as usize) {
// if A[i] is true:
if bool_vec[sieve_term] == true {
// for j = i^2, i^2+i, i^2+2i, i^2+3i, ...,
for j in SquareMultiple::new(sieve_term)
.take_while(|&term| term < max_val) { // ...not exceeding n
bool_vec[j] = false; // A[j] := false
}
}
}
let mut out_vec = vec![];
// Output: all i such that A[i] is true.
for term in 2..max_val {
if bool_vec[term] == true {
out_vec.push(term);
}
}
out_vec
}
Segmented Sieve
I struggled with this one the most. I wanted a .each_slice() method very badly, since that was the quickest route to the same algorithm in Ruby, and asked for it on SO.
/// A [segmented
/// approach](https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Segmented_sieve)
/// to sieveing, keeping memory use to O(√n). As Sorensen states, this is the
/// most practical optimization to the sieve of Eratosthenes.
///
/// # Examples
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
/// prime_suspects::segmented_sieve(65537, 256)
/// .iter().rev().take(3).map(|&num| num)
/// .collect::<Vec<usize>>());
/// ```
///
// ```
// assert_eq!(&999983,
// prime_suspects::segmented_sieve(1000000, 350).last().unwrap());
// ```
pub fn segmented_sieve(max_val: usize, mut segment_size: usize) -> Vec<usize> {
if max_val <= ((2 as i64).pow(16) as usize) {
// early return if the highest value is small enough (empirical)
return eratosthenes_sieve(max_val);
}
if segment_size > ((max_val as f64).sqrt() as usize) {
segment_size = (max_val as f64).sqrt() as usize;
println!("Segment size is larger than √{}. Reducing to {} to keep resource use down.",
max_val, segment_size);
}
// get the primes up to the first segment
let small_primes = eratosthenes_sieve((max_val as f64).sqrt() as usize);
let mut big_primes = small_primes.clone();
// As Sorensen says, we need to construct a sequence over each segment, in
// the interval [start + 1, start + segment_size] that begins with
// (start + this_prime - ( start mod p)), and increases by p up to
// (start + segment_size).
// That sequence will be the values to sieve out of this_segment.
// clunky way of doing each_slice, from
// https://stackoverflow.com/a/37033906/2023432
let mut segment_range = (segment_size..max_val).peekable();
while segment_range.peek().is_some() {
let this_segment: Vec<_> = segment_range.by_ref().take(segment_size).collect();
let mut sieved_segment: Vec<_> = this_segment.clone();
for &this_prime in &small_primes {
if !this_segment.is_empty() {
let mut starting_offset = this_segment[0] % this_prime;
starting_offset = if starting_offset == 0 { this_prime } else { starting_offset };
let first_val = this_segment[0] + this_prime - starting_offset;
let last_val: &usize = this_segment.last().unwrap();
// hack for inclusive range while RFC is figured out. see
// https://www.reddit.com/r/rust/comments/3xkfro/what_happened_to_inclusive_ranges/
let sieve_vec = (first_val..(*last_val + 1))
.step(this_prime)
.collect::<Vec<_>>();
sieved_segment = sieved_segment
.iter()
.filter(|&check_num| !sieve_vec.contains(&check_num))
.map(|&val| val)
.collect::<Vec<_>>();
}
}
for sieved_prime in sieved_segment {
big_primes.push(sieved_prime);
}
}
return big_primes;
}
#[test]
fn no_end_segment_sieve_misses() {
let test_100k_primes = segmented_sieve(100000, 300);
assert!(!test_100k_primes.contains(&99999));
let test_100m_primes = segmented_sieve(1000000, 350);
assert!(test_100m_primes.contains(&999983));
assert!(!test_100m_primes.contains(&999997));
}
