Generating prime numbers quickly in rust

Question

I am attempting to re-implement a postponed sieve algorithm for generating prime numbers in Rust. I am able to make a solution that works, but I have to use a couple of .clone() calls which I believe are killing my performance (the Rust solution ends up ~8x slower than the Python solution).

I would love some advice on how I can avoid the .clone() calls while avoiding errors from the borrow checker.

use std::collections::HashMap;

#[derive(Debug, Clone)]
struct Primes {
    i: usize,
    curr_candidate: u64,
    next_relevant_prime: u64,
    next_relevant_prime_squared: u64,
    sieve: HashMap<u64, u64>,
    initial_primes: Vec<u64>,
    internal_primes: Box<Option<Primes>>,
}

impl Primes {
    fn new() -> Primes {
        Primes {
            i: 0,
            curr_candidate: 7,
            next_relevant_prime: 0,
            next_relevant_prime_squared: 0,
            sieve: HashMap::new(),
            initial_primes: vec![2, 3, 5, 7],
            internal_primes: Box::new(None),
        }
    }
}

impl Iterator for Primes {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        let len = self.initial_primes.len();
        let mut internal_primes;
        if self.i < len {
            self.i += 1;
            return Some(self.initial_primes[self.i - 1]);
        } else if self.i == len {
            self.i += 1;
            internal_primes = Primes::new();
            self.internal_primes = Box::new(Some(internal_primes.clone()));
            internal_primes.next(); // skip 2
            self.next_relevant_prime = internal_primes.next().unwrap();
            self.next_relevant_prime_squared = self.next_relevant_prime.pow(2);
        } else {
            internal_primes = self.internal_primes.clone().unwrap();
        }
        let mut i = self.curr_candidate;
        loop {
            i += 2;
            let step;
            if self.sieve.contains_key(&i) {
                // composite
                step = self.sieve.remove(&i).unwrap();
            } else if i < self.next_relevant_prime_squared {
                // prime
                // save state for next round
                self.curr_candidate = i;
                self.internal_primes = Box::new(Some(internal_primes));
                return Some(i);
            } else {
                // i == next_relevant_prime_squared
                step = 2 * self.next_relevant_prime;
                self.next_relevant_prime = internal_primes.next().unwrap();
                self.next_relevant_prime_squared = self.next_relevant_prime.pow(2);
            }
            let mut j = i;
            j += step;
            while self.sieve.contains_key(&j) {
                j += step;
            }
            self.sieve.insert(j, step);
        }
    }
}

fn main() {
    let mut primes = Primes::new();
    for _i in 0..99_999 {
        primes.next();
    }
    println!("The 100,000th prime is {}", primes.next().unwrap())
}

Answer 1

I was able to avoid the calls to .clone() with two changes.

First, we can avoid setting

self.internal_primes = Box::new(Some(internal_primes.clone()));

at all, and instead delay assigning internal_primes to self.internal_primes until right before the function returns. This avoids having multiple mutable references to the same Primes instance.

Second, we can use the other .clone() call by using mem::replace:

internal_primes = mem::replace(&mut self.internal_primes, Box::new(None)).unwrap();

I also realized that cargo run is running in debug mode (at least for my project), which is much slower than I would have though. Compiling for release mode ends up being much faster regardless of whether we use .clone() (although it's still about 2x faster to avoid them):

$ cargo build --release
$ time ./target/release/primes
The 100,000th prime is 1299709
./target/release/primes  0.07s user 0.00s system 34% cpu 0.216 total

The final version of the code:

use std::collections::HashMap;
use std::mem;

#[derive(Debug, Clone)]
struct Primes {
    i: usize,
    curr_candidate: u64,
    next_relevant_prime: u64,
    next_relevant_prime_squared: u64,
    sieve: HashMap<u64, u64>,
    initial_primes: Vec<u64>,
    internal_primes: Box<Option<Primes>>,
}

impl Primes {
    fn new() -> Primes {
        Primes {
            i: 0,
            curr_candidate: 7,
            next_relevant_prime: 0,
            next_relevant_prime_squared: 0,
            sieve: HashMap::new(),
            initial_primes: vec![2, 3, 5, 7],
            internal_primes: Box::new(None),
        }
    }
}

impl Iterator for Primes {
    type Item = u64;

    fn next(&mut self) -> Option<Self::Item> {
        let len = self.initial_primes.len();
        let mut internal_primes;
        if self.i < len {
            self.i += 1;
            return Some(self.initial_primes[self.i - 1]);
        } else if self.i == len {
            self.i += 1;
            internal_primes = Primes::new();
            internal_primes.next(); // skip 2
            self.next_relevant_prime = internal_primes.next().unwrap();
            self.next_relevant_prime_squared = self.next_relevant_prime.pow(2);
        } else {
            internal_primes = mem::replace(&mut self.internal_primes, Box::new(None)).unwrap();
        }
        let mut i = self.curr_candidate;
        loop {
            i += 2;
            let step;
            if self.sieve.contains_key(&i) {
                // composite
                step = self.sieve.remove(&i).unwrap();
            } else if i < self.next_relevant_prime_squared {
                // prime
                // save state for next round
                self.curr_candidate = i;
                self.internal_primes = Box::new(Some(internal_primes));
                return Some(i);
            } else {
                // i == next_relevant_prime_squared
                step = 2 * self.next_relevant_prime;
                self.next_relevant_prime = internal_primes.next().unwrap();
                self.next_relevant_prime_squared = self.next_relevant_prime.pow(2);
            }
            let mut j = i;
            j += step;
            while self.sieve.contains_key(&j) {
                j += step;
            }
            self.sieve.insert(j, step);
        }
    }
}

fn main() {
    let mut primes = Primes::new();
    for _i in 0..99_999 {
        primes.next();
    }
    println!("The 100,000th prime is {}", primes.next().unwrap())
}

Answer 2

A normal Sieve of Eratosthenes would look like:

fn primes(n: usize) -> impl Iterator<Item = usize> {
    const START: usize = 2;
    if n < START {
        Vec::new()
    } else {
        let mut is_prime = vec![true; n + 1 - START];
        let limit = (n as f64).sqrt() as usize;
        for i in START..limit + 1 {
            let mut it = is_prime[i - START..].iter_mut().step_by(i);
            if let Some(true) = it.next() {
                it.for_each(|x| *x = false);
            }
        }
        is_prime
    }
    .into_iter()
    .enumerate()
    .filter_map(|(e, b)| if b { Some(e + START) } else { None })
}

1. Starting at an offset of 2 means that an n < 2 input requires zero allocations because Vec::new() doesn't allocate memory until elements are pushed into it.
2. Using Vec as an output to the if .. {} else {} condition means the output is statically deterministic, avoiding the need for a boxed trait object.
3. Iterating is_prime with .iter_mut() and then using .step_by(i) makes all the optimizations required, and removes a lot of tediousness.
4. Returning impl Iterator allows for static dispatching instead of dynamic dispatching, which is possible because the type is now statically known at compile-time, making the zero input/output condition order of magnitude faster.

If you need to calculate primes as the range increases above 10's of millions then check out https://rosettacode.org/wiki/Sieve_of_Eratosthenes#Unbounded_Page-Segmented_bit-packed_odds-only_version_with_Iterator