Sieve of Eratosthenes in Rust

Question

This is a relatively simple implementation of the Sieve of Eratosthenes in Rust. The main objective is to find the \$n\$th prime quickly when \$n\$ might grow to huge numbers.

pub fn nth(n: u32) -> u32 {
    // A convenient variable that will help in simplifying a complicated expression
    let x = if n <= 10 { 10.0 } else { n as f64 };

    // The prime counting function is pi(x) which is approximately x/ln(x)
    // We need the inverse of that. A good estimate is ceil(x * ln(x * ln(x)))
    let limit: usize = (x * (x * (x).ln()).ln()).ceil() as usize;
    let mut sieve = vec![true; limit];
    let mut count = 0;

    // Exceptional case for 0 and 1
    sieve[0] = false;
    sieve[1] = false;

    for prime in 2..limit {
        if !sieve[prime] {
            continue;
        }
        if count == n {
            return prime as u32;
        }
        count += 1;

        for multiple in ((2 * prime)..limit).step_by(prime) {
            sieve[multiple] = false;
        }
    }
    return <u32>::max_value();
}

This entire function would be present in main.rs or lib.rs, as necessary.

Please feel free to nitpick on any aspect of the code. Suggestions to make the code cleaner or more efficient are welcome.

Answer 1

Nice use of the Prime Number Theorem there. However, there are some issues with your code as is:

• The name nth does not convey the meaning of the function.
• A return statement is used instead of an expression at the end of the function.
• The composite marking can start at prime * prime, as all smaller composites are already marked by previous primes.
• The \$\left(2^{32}-1\right)\text{th}\$ prime is larger than \$2^{32}-1\$, but less than \$2^{64}-1\$, so u64 is a better return type.
• The function will use a lot of memory if we use <u32>::max_value() as input.
• The first (\$1\text{st}\$) prime number \$p_1\$ is \$2\$, but your code returns \$3\$ on nth(1).
• The behaviour is not documented at the moment, as the function is missing all documentation.
• The function doesn't necessarily return a prime number.

The last issue is the greatest one, to be honest. If I use a nth_prime function, I expect the function to return a prime number and only a prime number, whereas <u32>::max_value() is a composite. If it's possible that your function does not return a valid result, use Option or Result instead:

// https://play.rust-lang.org/?gist=4d6abc78a8c0d205da57a17c02201d7c&version=stable&mode=release&edition=2015

/// Returns the nth prime.
///
/// It uses a sieve internally, with a size of roughly
/// `n * (n.ln() + n.ln().ln()` bytes. As a result, its
/// runtime is also bound loglinear by the upper term.
///
/// # Examples
///
/// ```
/// use <yourcrate>::nth_prime;
///
/// assert_eq!(nth_prime(0), None);
/// assert_eq!(nth_prime(1), Some(2));
/// assert_eq!(nth_prime(2), Some(3));
/// assert_eq!(nth_prime(3), Some(5));
/// ```
///
/// 
pub fn nth_prime(n: u32) -> Option<u64> {
    if n < 1 {
        return None;
    }

    // The prime counting function is pi(x) which is approximately x/ln(x)
    // A good upper bound for the nth prime is ceil(x * ln(x * ln(x)))
    let x = if n <= 10 { 10.0 } else { n as f64 };
    let limit: usize = (x * (x * (x).ln()).ln()).ceil() as usize;
    let mut sieve = vec![true; limit];
    let mut count = 0;

    // Exceptional case for 0 and 1
    sieve[0] = false;
    sieve[1] = false;

    for prime in 2..limit {
        if !sieve[prime] {
            continue;
        }
        count += 1;
        if count == n {
            return Some(prime as u64);
        }

        for multiple in ((prime * prime)..limit).step_by(prime) {
            sieve[multiple] = false;
        }
    }
    None
}

Note that the documentation's examples show that 0 isn't a viable index for primes, and that nth_prime(1) == Some(2).

Other than that, we can optimize the code a little bit further if we consider 2 as a special case and then step (3..limit).step_by(2), but that makes the code slightly harder to read. There are also several approaches to keep the memory usage limited, but they are harder to achieve.

Other than that, well done, although I'd link to the Prime Number Theorem for the approximation.

Answer 2

    // Exceptional case for 0 and 1
    sieve[0] = false;
    sieve[1] = false;

KISS. Deleting these lines of code will not affect correctness in any way, so delete them to remove distractions.

---

    for prime in 2..limit {
        if !sieve[prime] {
            continue;
        }

sieve is not the most helpful name, because some people write their sieves with true indicating a composite number and others with true indicating a prime number. is_prime would be explicit. If !is_prime[prime] offends your aesthetics then prime could also be renamed.

---

        if count == n {
            return prime as u32;
        }

The need to cast here is worrying, and suggests that the types are wrong: either limit should be u32 or the return value should be u64.

---

The main objective is to find the \$n\$th prime quickly when \$n\$ might grow to huge numbers.

What counts as huge? Depending on the answer, it might make sense to trade increased complexity for a reduction in memory usage and segment the sieve. It might even be worth pre-calculating (or borrowing from resources like the Prime Pages) a lookup table of every millionth prime and switching to a sieve which allows you to sieve an arbitrary range rather than having to start at 2.