Generators - Prime sieve

I miss Python's generators in (stable) Rust.

Here is an example where I miss it - an incremental prime sieve implemented in py3 and rust.
The rust version is ~70% longer (101 lines vs 61).

How would you have done the rust version?

1 Like

Using impl Iterator return types, this becomes shorter.

/// Sorenson's rolling sieve - yield next prime each time it is called.
/// Complexity: O(nloglogn) time, O(sqrt(n)log(n)) bits,
/// O(logn/loglogn) incremental   
/// See algorithm description in paper:
/// "Two Compact Incremental Prime Sieves", Jonathan P. Sorenson, Journal of Computation and Mathematics, 2015
/// https://arxiv.org/abs/1503.02592
fn rolling_sorenson() -> impl Iterator<Item = u64> {
    let start = 100;

    let primes = [
        2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89,
        97,
    ];

    let mut r = f64::sqrt(start as f64) as u64 + 1;
    let mut s = r * r;
    let mut delta = r + 2;
    let mut t = vec![vec![]; (delta + 1) as usize];
    for p in primes.into_iter().take_while(|&p| p < r - 1) {
        let j = (p - (start % p)) % p;
        t[j as usize].push(p);
    }
    let mut pos = 0;
    let mut n = start;

    primes.into_iter().chain(
        std::iter::repeat_with(move || {
            let mut is_prime = true;
            while let Some(p) = t[pos as usize].pop() {
                t[((pos + p) % delta) as usize].push(p);
                is_prime = false;
            }
            if n == s {
                if is_prime {
                    t[((pos + r) % delta) as usize].push(r);
                    is_prime = false;
                }
                r += 1;
                s = r * r;
            }
            n += 1;
            pos = (pos + 1) % delta;
            if pos == 0 {
                delta += 2;
                t.extend([vec![], vec![]]);
            }
            is_prime.then_some(n - 1)
        })
        .flatten(),
    )
}

use clap::Parser;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    #[arg(short, long, default_value_t = 100)]
    ///primes below
    n: u64,
}

fn main() {
    let args = Args::parse();

    let mut rs = rolling_sorenson();
    loop {
        let p = rs.next().unwrap();
        if p >= args.n {
            break;
        }
        println!("{}", p);
    }
}

(playground)

There can probably be done more to this code to make it even more rusty.

I conservatively used u64 everywhere; I haven’t yet tried to identify which variables could be made usize instead.

First of all, number of lines is a terrible measure of code size or complexity. It's trivially influenced by how liberally you use empty lines, for example.

Second, Rust is a curly brace language, so you'll have 1 extra line for each scope, that in itself accounts for 17 lines for this code. That doesn't really increase conciseness, readability, or complexity of the code at all. There are also lines in the Python code that contract a normally multi-line statement into one, e.g. if isPrime: yield n-1

There is also no need for the VecDeque. You are only using the initial vector of primes to return them one by one. Vec::into_iter() or Vec::drain() can do that just fine. So you could just .chain() the actual implementation onto the end of the set of hard-coded seeding primes.

With the suggested modifications, the Rust code becomes 86 lines, and it's not any more convoluted than the Python version; the main source of it being longer is now the declaration and instantiation of the RollingSorensen struct, which is 2 * 8 = 16 lines, accounting for almost all of the code size difference.

TL;DR: the Rust code being somewhat longer seems to have nothing to do with generators.

3 Likes

One small thing: I guess, the let n = start; std::iter::repeat_with(move || { …; n += 1; …; is_prime.then(n-1) }.flatten() I came up with above is more naturally written as (start..).filter(move |&n| { …; …; is_prime }). So the thing looks like

fn rolling_sorenson() -> impl Iterator<Item = u64> {
    let start = 100;

    let primes = [
        2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89,
        97,
    ];

    let mut r = f64::sqrt(start as f64) as u64 + 1;
    let mut s = r * r;
    let mut delta = r + 2;
    let mut t = vec![vec![]; (delta + 1) as usize];
    for p in primes.into_iter().take_while(|&p| p < r - 1) {
        let j = (p - (start % p)) % p;
        t[j as usize].push(p);
    }
    let mut pos = 0;

    primes.into_iter().chain((start..).filter(move |&n| {
        let mut is_prime = true;
        while let Some(p) = t[pos as usize].pop() {
            t[((pos + p) % delta) as usize].push(p);
            is_prime = false;
        }
        if n == s {
            if is_prime {
                t[((pos + r) % delta) as usize].push(r);
                is_prime = false;
            }
            r += 1;
            s = r * r;
        }
        pos = (pos + 1) % delta;
        if pos == 0 {
            delta += 2;
            t.extend([vec![], vec![]]);
        }
        is_prime
    }))
}

Edit: Also, my main function above can be simplified to

fn main() {
    let args = Args::parse();

    for p in rolling_sorenson().take_while(|&p| p < args.n) {
        println!("{}", p);
    }
}

The Rust program is already down to 64 lines now.

5 Likes

Cool - I like this solution. I picked up several things from you here, but main one is returning an iterator, because this avoids the need for a Self struct to keep the state - which is also what the generator does in the py3 version.