Implement Haskell-like primes iterator in Rust

user_5600363 0
Viewed 148

I am trying to implement a Rust iterator similar to the following Haskell code:

primes :: [Integer]
primes = nextPrime [2..]
    where 
        nextPrime (x:xs) = x : nextPrime (filter (notDivBy x) xs)
        notDivBy a x = a `mod` x /= 0

My attempt so far (playground):

// as pointed out in a comment this can simply return Primes and not use Box::new
fn primes() -> Box<dyn Iterator<Item = usize>> {
    Box::new(Primes::new())
}

struct Primes {
    nums: Box<dyn Iterator<Item = usize>>,
}

impl Primes {
    fn new() -> Self {
        Primes { nums : Box::new(2..) }
    }
}

impl Iterator for Primes {
    type Item = usize;
    
    fn next(&mut self) -> Option<usize> {
        
        let prime = self.nums.next().unwrap();
        self.nums = Box::new(self.nums.filter(move |&n|!divides(prime, n)));
        //use std::borrow::BorrowMut;
        //*self.nums.borrow_mut() = self.nums.filter(move |&n|!divides(prime, n));
        Some(prime)
    }
}


pub fn divides(d: usize, n: usize) -> bool {
    n % d == 0
}

Unfortunately this runs into:

error[E0507]: cannot move out of `self.nums` which is behind a mutable reference
  --> src/lib.rs:22:30
   |
22 |         self.nums = Box::new(self.nums.filter(move |&n| !divides(prime, n)));
   |                              ^^^^^^^^^ move occurs because `self.nums` has type `Box<dyn Iterator<Item = usize>>`, which does not implement the `Copy` trait

For more information about this error, try `rustc --explain E0507`.

Or if you uncomment the alternative borrow_mut code:

error[E0277]: the trait bound `Box<(dyn Iterator<Item = usize> + 'static)>: BorrowMut<Filter<Box<dyn Iterator<Item = usize>>, [closure@src/lib.rs:24:52: 24:79]>>` is not satisfied
  --> src/lib.rs:24:20
   |
24 |         *self.nums.borrow_mut() = self.nums.filter(move |&n|!divides(prime, n));
   |                    ^^^^^^^^^^ the trait `BorrowMut<Filter<Box<dyn Iterator<Item = usize>>, [closure@src/lib.rs:24:52: 24:79]>>` is not implemented for `Box<(dyn Iterator<Item = usize> + 'static)>`
   |
   = help: the following implementations were found:
             <Box<T, A> as BorrowMut<T>>

Frankly I am not even sure which of these two is closer to working.

rust primes iterator haskell
3 Answers

I think you can cache the discovered primes:

fn primes() -> Primes {
    Primes::new()
}

struct Primes {
    nums: Box<dyn Iterator<Item = usize>>,
    cache: Vec<usize>,
}

impl Primes {
    fn new() -> Self {
        Primes {
            nums: Box::new(2..),
            cache: Vec::new(),
        }
    }
}

impl Iterator for Primes {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        loop {
            let num = self.nums.next().unwrap();
            if self
                .cache
                .iter()
                .take_while(|&&n| n*n <= num)
                .all(|&p| !divides(p, num))
            {
                // we found a prime!
                self.cache.push(num);
                return Some(num);
            }
        }
    }
}

pub fn divides(d: usize, n: usize) -> bool {
    n % d == 0
}

fn main() {
    println!("Primes: {:?}", primes().take(10).collect::<Vec<_>>());
}

Playground

user_16951720

I'm not really expert in Rust, but one (ugly) option is to replace self.nums with a dummy iterator so to be able to move the previous value.

fn next(&mut self) -> Option<usize> {
    let prime = self.nums.next().unwrap();
    // Ugly replacement with a dummy iterator value
    let rest = std::mem::replace(&mut self.nums, Box::new(0..));
    self.nums = Box::new(rest.filter(move |&n|!divides(prime, n)));
    Some(prime)
}

This could be cleaner if we used an Option wrapper so that we can use None as the dummy value.

user_32349590

I finally figured out (with the help of the other answers) how to fix cannot move out of self.nums which is behind a mutable reference. The trick is to use &mut self.nums (and get rid of the Box).

fn primes() -> Primes {
    Primes::new()
}

pub struct Primes {
    nums: std::ops::RangeFrom<usize>,
    cache: Vec<usize>,
}

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

impl Iterator for Primes {
    type Item = usize;

    fn next(&mut self) -> Option<usize> {
        let next_prime = (&mut self.nums)
            .filter(|&n| {
                self.cache
                    .iter()
                    .take_while(|&p| p * p <= n)
                    .all(|&p| !divides(p, n))
            })
            .next()
            .unwrap();
        self.cache.push(next_prime);
        return Some(next_prime);
    }
}

pub fn divides(d: usize, n: usize) -> bool {
    n % d == 0
}

fn main() {
    println!("Primes: {:?}", primes().take(10).collect::<Vec<_>>());
}

playground

user_56003630
Related