Prime sieve in Rust

Question

Using Sieve of Eratosthenes, I created a function that returns a vector of primes up to given limit.

fn prime_until(limit: usize) -> Option<Vec<usize>> {
    if limit <= 1 {
        return None;
    }

    // limit + 1 because of zero based indexing
    let mut nums = vec![true; limit + 1];
    nums[0] = false; // zero is not a prime
    nums[1] = false;

    for i in 0..nums.len() {
        if !nums[i] {
            continue;
        }

        let mut j = i * i;
        while j <= limit {
            nums[j] = false;
            // lets say j == i * 2,
            // j + i == i * 3
            // basically increment the multiplication
            j += i;
        }
    }

    Some(
        nums.iter()
            .enumerate()
            .filter(|(_, is_prime)| **is_prime)
            .map(|(i, _)| i)
            .collect(),
    )
}

On my computer, it can produce the primes up to 100,000,000 (5,761,455 primes) in ~1.8s when compiling with cargo build --release.

$ time ./sieve_of_eratosthenes
[src/main.rs:8] &primes.as_ref().unwrap().len() = 5761455

real    0m1.698s
user    0m1.628s
sys     0m0.067s

Update #1: Using this line seems to speed things up a little bit. Though I heard that there might be some problem with precision etc.

...
    nums[0] = false; // zero is not a prime
    nums[1] = false;

    for i in 0..f32::sqrt(nums.len() as f32) as usize { // here
        if !nums[i] {
            continue;
        }
...

$ time ./sieve_of_eratosthenes
[src/main.rs:8] &primes.as_ref().unwrap().len() = 5761455

real    0m1.554s
user    0m1.504s
sys     0m0.048s

Are my implementation good? What can be improved?

Answer 1

Why return an Option<Vec> when all the None signals is that there are no elements, just returning an empty vector (Vec::new()) tells the same story.

Then since you don't use nums afterwards you can use into_iter instead of iter to get rid of one level of indirection.

The use of .filter().map() always looks like a missed opportunity to use filter_map getting rid of the last indiretion in that iteration.

    nums.into_iter()
        .enumerate()
        .filter_map(|(i, is_prime)| is_prime.then_some(i))
        .collect()

On the naming part of things nums isn't that descriptive of a name, especially since you're not actually holding numbers, maybe is_prime is more descriptive, though I'm definitely no good at coming up with names either. But with the new name this comment:

    is_prime[0] = false; // zero is not a prime

looks rather silly, in my opinion it's rather clear what is_prime[0] = false; means even without it.

I dislike handling of indices manually so I would write the inner while loop as this functional iterator chain instead:

        is_prime.iter_mut()
            .skip(i * i) // start at i*i since smaller multiples have been removed earlier
            .step_by(i) // only visit multiples of i
            .for_each(|p| *p = false);

Starting at 0 when you already explicitly initialized the first 2 values is also pretty redundant so the outer loop can start at 2 instead.
You should be able to gain the speed boost of stopping early while still being accurate by simply squaring both sides instead of taking the root, both these together give an outer loop of:

for i in 2.. {
    if i * i >= is_prime.len() {
        break;
    }
    // ...
}

You could also invert the logic to leverage the fact that usually memory pages allocated from the operating system are zeroed from the get go so no instructions have to be added to set them to false.

All of this together gives us this listing:

fn prime_until(limit: usize) -> Vec<usize> {
    if limit <= 1 {
        return vec![];
    }

    // limit + 1 because of zero based indexing
    let mut not_prime = vec![false; limit + 1];
    not_prime[0] = true;
    not_prime[1] = true;

    for i in 2usize.. {
        let j = i.saturating_mul(i);
        if j >= not_prime.len() {
            break;
        }
        if not_prime[i] {
            continue;
        }
        
        not_prime.iter_mut()
            .skip(j) // start at i*i since smaller multiples have been marked earlier
            .step_by(i) // only visit multiples of i
            .for_each(|n| *n = true);
    }

    not_prime.into_iter()
        .enumerate()
        .filter_map(|(i, not_prime)| (!not_prime).then_some(i))
        .collect()
}