Rust version of basic functional prime stream - can this be made less clunky?

Question

I'm looking at implementing the basic functional not-quite-Eratosthenes prime stream in Rust. I like to try it when I start learning a language.

Here's the bog standard Haskell version:

primes :: [Integer]
primes = sieve (2 : [3, 5..])
  where
    sieve (p:xs) = p : sieve [x|x <- xs, x `mod` p /= 0]

in Rust, I think a reasonable transliteration is something like:

struct Primes {
    iter: Option<Box<Iterator<Item=u32>>>
}

impl Iterator for Primes {
    type Item = u32;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        let mut iter = self.iter.take().unwrap();
        let res = iter.next().unwrap();
        self.iter = Some(Box::new(iter.filter( move |x| x % res != 0)));
        Some(res)
    }
}

fn main() {
    let primes = Primes { iter: Some(Box::new(2..)) };
    for p in primes.take(20) {
        println!("{}", p);
    }
}

Leaving aside the algorithm (of course there are nicer ways to do primes!), this feels a bit verbose - is there a simpler way to construct an iterator with accumulating filters? The Option of Box feels quite clunky, but if I understand right, that's the best way to replace a field on self.

Answer 1

On a first note, your code will panic once your iterator reaches std::u32::MAX, which admittedly would take a long time. You can fix this by returning None when this happens, which can be done easily with ?.

let mut iter = self.iter.take()?;
let res = iter.next()?;

I can't think of a better way to accumulate actual filter calls, but if I were to implement this sort of algorithm, I would probably do something like this instead:

struct Primes {
    found: Vec<u32>,
    iter: std::ops::RangeFrom<u32>,
}

impl Primes {
    fn new() -> Self {
        Primes {
            found: Vec::new(),
            iter: (2..),
        }
    }
}

impl Iterator for Primes {
    type Item = u32;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        let Primes { found, iter } = self;
        let res = iter.find(|x| found.iter().all(|p| x % p != 0))?;
        found.push(res);
        Some(res)
    }
}

Note that in your code, all the previous primes are stored in the closures, so doing it explicitly this way gives a better view of the costs of this method. Additionally, I can imagine the continuously nested dynamic dispatch adding up in costs, so this method removes any need for that.