1
\$\begingroup\$

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
}
\$\endgroup\$

1 Answer 1

3
\$\begingroup\$

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.

\$\endgroup\$
8
  • \$\begingroup\$ Thanks for the tips. I especially appreciate the comments on the sequence of calls at the end. Regarding this line: for p in prime_indices.iter_mut().skip(i * i).step(2 * i) { Why does step(2 * i) work? Is it b/c of the step where you've preempted the multiples of 2? Interestingly for value of x up to 10000, this runs about the same to a little slower on my machine. 0.20 s vs 0.22 s on average. Probably just noise since I didn't run any true benchmarks. I wasn't aware there was a version of rustfmt that ran on stable. Good to know. I'll definitely make use of it going forward. \$\endgroup\$ Sep 7, 2017 at 2:16
  • \$\begingroup\$ (1) step(2 * i) is used instead of step(i) in order to traverse odd multiples only, as mentioned in the first paragraph. (2) That is odd, did you build for the Release target? (3) rustfmt hasn't reached 1.0 yet, but it can be installed with a stable toolchain and works with stable code. \$\endgroup\$ Sep 7, 2017 at 9:40
  • \$\begingroup\$ Minor clippy thing: u32 cannot be < 0: just compare == 0 instead of <= 0. Read more \$\endgroup\$
    – CAD97
    Sep 7, 2017 at 23:50
  • \$\begingroup\$ @CAD97 Good point, somehow I had the illusion that i was dealing with a signed integer. Thanks, will update. \$\endgroup\$ Sep 7, 2017 at 23:54
  • 1
    \$\begingroup\$ @MatthewStevenson the rustfmt-nightly crate is the most updated version of rustfmt, and only runs on the nightly toolchain. However, if you've installed the nightly toolchain rustup install nightly, you can install rustfmt cargo +nightly install rustfmt-nightly and use it cargo +nightly fmt on stable code even without changing your toolchain. Or just use the rustfmt crate which is an old no-longer-updated version that works on stable. Eventually the goal is it'll be distributed with the stable compiler and thus work on stable. \$\endgroup\$
    – CAD97
    Sep 7, 2017 at 23:54

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.