Incremental Sieve of Eratosthenes using iterators

Question

I am in the process of learning Rust and am still having some issues with borrowing and ownership. I find myself trying to borrow mutable references that are already borrowed as immutable references etc...

I noticed that when I try to iterate over the primes without owning them, I encounter issues trying to modify the composites hash map. I was able to get around this by calling to_owned on the vector. Is this the right way to handle this?

I also see that when I build the code, I get warnings about unused code for warning: struct is never constructed: 'Sieve' and warning: method is never used: 'new', am I constructing them incorrectly?

use std::collections::HashMap;

struct Sieve {
    composites: HashMap<u64, Vec<u64>>,
    x: u64,
}

impl Sieve {
    fn new() -> Sieve {
        Sieve {
            composites: HashMap::new(),
            x: 2,
        }
    }
}

impl Iterator for Sieve {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        let x = self.x;
        self.x = self.x + 1;
        match self.composites.get(&x) {
            Some(numbers) => {
                for _num in numbers.to_owned() {
                    self.composites
                        .entry(x + _num)
                        .and_modify(|v| v.push(_num))
                        .or_insert(vec![_num]);
                }
                self.composites.remove(&x);
                self.next()
            }
            None => {
                self.composites.insert(x * x, vec![x]);
                Some(x)
            }
        }
    }
}

fn main() {
    let mut sieve = Sieve::new();

    println!("{:?}", sieve.next()); // 2
    println!("{:?}", sieve.next()); // 3
    println!("{:?}", sieve.next()); // 5
    println!("{:?}", sieve.next()); // 7
}

Here is the code on the Rust playground.

I previously posted a version of the Sieve of Eratosthenes using experimental Rust features, this made it difficult to get feedback on. I have refactored the code to use iterators

Answer 1

Your code is totally fine and idiomatic, but I would change some minor points:

use std::collections::HashMap;

struct Sieve {
    composites: HashMap<u64, Vec<u64>>,
    current: u64,
}

impl Default for Sieve {
    fn default() -> Self {
        Sieve {
            composites: HashMap::new(),
            current: 2,
        }
    }
}

impl Sieve {
    pub fn new() -> Sieve {
        Default::default()
    }
}

impl Iterator for Sieve {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        fn next_prime(composites: &mut HashMap<u64, Vec<u64>>, x: u64) -> u64 {
            match composites.get(&x) {
                Some(numbers) => {
                    for num in numbers.to_owned() {
                        composites
                            .entry(x + num)
                            .and_modify(|v| v.push(num))
                            .or_insert_with(|| vec![num]);
                    }
                    composites.remove(&x);
                    next_prime(composites, x + 1)
                }
                None => {
                    composites.insert(x * x, vec![x]);
                    x
                }
            }
        }

        let prime = next_prime(&mut self.composites, self.current);
        self.current = prime + 1; // This number will be the next to be tested

        Some(prime)
    }
}

fn main() {
    let mut sieve = Sieve::new();

    assert_eq!(sieve.next(), Some(2));
    assert_eq!(sieve.next(), Some(3));
    assert_eq!(sieve.next(), Some(5));
    assert_eq!(sieve.next(), Some(7));
}

• It is good to implement Default when you can build your data structure without parameters (related answer).
• This is a good practice to put assertions instead of prints. That's easier to change to code and verify that it is still ok.
• The naming of the variables is important; but I do not pretend that mine is perfect, though. By the way, _num should not be prefixed with an underscore because it is used.
• You can run clippy to catch some common error. Clippy warns you that you create a vec![_num] at each iteration. You should give to or_insert_with a closure that will build the correct vector.
• The more important refactoring that I did is to make the algorithm more "functional". I think that yours is difficult to reason about because it calls recursively the Iterator::next method and rely on internal state.

Answer 2

I have some more suggestions that can improve your code.

• You can avoid cloning numbers by just removing it from composites to begin with, since you already remove it later. HashMap::remove will give you an Option<Vec<u64>>, so you can now modify the hashmap while you iterate over the vec because it is no longer owned by the hashmap.
• You can make your function iterative by using a while let loop.
• You can simplify your usage of entry by just using or_default which will give a mutable reference to either the vec that was already there, or an empty one. Then you can just push into that vec.
• Another option for clarity is to make a next_prime method on Sieve.

impl Sieve {
    pub fn next_prime(&mut self) -> u64 {
        while let Some(numbers) = self.composites.remove(&self.current) {
            for num in numbers {
                self.composites
                    .entry(self.current + num)
                    .or_default()
                    .push(num)
            }
            self.current += 1;
        }
        let prime = self.current;
        self.composites.insert(prime * prime, vec![prime]);
        self.current += 1;
        prime
    }
}

impl Iterator for Sieve {
    type Item = u64;

    fn next(&mut self) -> Option<u64> {
        Some(self.next_prime())
    }
}