Rust nth prime implementation

Question

I have implemented an nth prime algorithm in Rust that uses value lookup for n < 6, trial division for n < 1000, and a sieve function for n > 999. How can it be better optimized?

pub fn nth(x: u32) -> Result<u64, String> {

match x > 0 {
    true => match x {
        1 => Ok(2),
        2 => Ok(3),
        3 => Ok(5),
        4 => Ok(7),
        5 => Ok(11),
        x if x < 1000 => Ok(trial_div(x)),
        _ => Ok(sieve(x)),
    },
    false => Err(String::from("invalid input")),
  }
}

fn trial_div(x: u32) -> u64 {
  let mut primes: Vec<u64> = Vec::with_capacity(x as usize);

  primes.push(2);

  if x > 1 {
    primes.push(3);
  } 

  let mut next_checked = *primes.last().unwrap() + 2;

  while primes.len() < x as usize {
        if primes.iter().all(|&i| next_checked % i != 0) {
            primes.push(next_checked)
        }
        next_checked += 2;
  }
  *primes.last().unwrap()
}

fn sieve(x: u32) -> u64 {
  let x = x as f64;

  let upper_limit = (x * ((x*x.ln()).ln())).floor() as usize;
  let mut prime_indices = vec![true; upper_limit];

  for i in 2..upper_limit {
    if prime_indices[i] {
        let mut counter = 0;
        while (i.pow(2) + (counter * i)) < upper_limit {
            prime_indices[(i.pow(2) + (counter * i)) as usize] = false;
            counter += 1
        }
    }
  }

  let mut prime_x = 1;
  let mut prime_count = 0;
  for i in 2..upper_limit {
    if prime_indices[i] {
        prime_count += 1;
        prime_x = i
    }
    if prime_count == x as u64 {
        break
    }
  }
  prime_x as u64
}

Answer 1

There are some optimizations from the original algorithm that can be included here: the base iteration of the prime p before the sieve is filled with its multiples can start with p = 2 and subsequently with odd numbers only (2, 3, 5, ..., upper_limit). Another useful refinement is that starting from 3, we only need to set the indices of the sieve with a step by 2 p. So for p = 3, we'd only have to clear (9, 15, 21, 27, ...). The rest of the review will be focused around idiomatic Rust, while integrating these optimizations.

In function nth, matching on a boolean value is something awkward to do, since if statements are also expressions and are more readable. Even better here is to turn it into another match clause for the error. Values of type u32 cannot be negative, so testing for 0 is enough.

match x {
    0 => Err("invalid input".to_string()),
    1 => Ok(2),
    2 => Ok(3),
    3 => Ok(5),
    4 => Ok(7),
    5 => Ok(11),
    x if x < 1000 => Ok(trial_div(x)),
    x => Ok(sieve(x)),
}

To create a String out of a &str, you can also use the to_string() or to_owned() methods. "invalid input".into() would also work here because the program already expects a String for the error type and String implements From<&str>.

We can see multiple uses of indexing operations over the prime_indices vector (e.g. prime_indices[i.pow(2) + (counter * i)] and if prime_indices[i]). While these uses are not necessarily wrong, the operation is bound-checked at run-time and it isn't guaranteed that the compiler will optimize them out. The preferred way of manipulating collections and other sequences of data is with the Iterator API. The use of iterators can become tricky when dealing with mutable collections, since there cannot be multiple borrows to the same vector. Nevertheless, we can remove most of these uses, resulting in better optimized and more idiomatic code, potentially less prone to indexing mistakes that could lead to a panic.

At the beginning of sieve, line 42, we have one of these cases where it is best to leave the ranged iteration alone, because the inner loop will modify the same content the outer loop is attempting to read. When filling the sieve for each multiple of i, we can rethink this operation over the vector as a stepped iteration starting from i.pow(2) (this can be replaced with i * i) with a step i. The step_by iteration adaptor is unfortunately not stable yet, but there is an equivalent Itertools#step method in the itertools crate. The result:

for i in 2..upper_limit {
    if prime_indices[i] {
        for p in prime_indices.iter_mut().skip(i * i).step(i) {
            *p = false;
        }
    }
}

Or even better, with the aforementioned optimizations:

for p in prime_indices.iter_mut().skip(4).step(2) {
    *p = false;
}
for i in (3..upper_limit).step(2) {
    if prime_indices[i] {
        for p in prime_indices.iter_mut().skip(i * i).step(2 * i) {
            *p = false;
        }
    }
}

Then, we want the xth prime number from the vector. We can achieve the same outcome as your loop with the following sequence of adaptors, commented one by one for clarity:

let prime_x = prime_indices.into_iter()
    .enumerate()           // keep track of original indices
    .skip(2)               // start at number 2
    .filter(|&(_, p)| p)   // primes only
    .map(|(i, _)| i)       // discard boolean
    .nth((x - 1) as usize) // take the xth prime (1-indexed minus 2 skipped)
    .unwrap_or(1);         // default to 1

My final advice is to consider running the rustfmt tool to format the entire code, as I've seen some irregular indentations. Here is the full result in a Rust Playground.