nth-prime.rs

/*

Incremental sieve of Erastotenes
================================

We get the nth prime by an variation of 'sieve of Erastotenes' that is
incremental and does not require a upper bound.

You can visualize the algorithm as it follows. Imagine some boxes
labeled with the integer numbers.


   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
   ^
   i


i points to the current integer candidate. We will move it to the
right for each iteration.

The box for 2 is empty, so it is prime. We add it to our list of
primes and mark its square (4), so it is considered composite.

   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
   ^         (2*)
   i

Continue


   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
         ^   (2*)
         i


   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
         ^   (2*)                           (3*)
         i

  2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
             (2*)                           (3*)
               ^
               i

We have reached 4. But there is a (2*) mark in there. So we know is
composite. Then move the (2*) mark to the next multiple of the prime.

   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
               ^          (2*)             (3*)
               i

and we continue to find the next primes.

   2 --- 3 --- 4 --- 5 --- 6 --- 7 --- 8 --- 9
                     ^    (2*)             (3*)
                     i
 */

pub fn nth(n: u32) -> u32 {
    // Hold the next multiple of the primes, like (2,4), (3,12), ...
    let mut sieve: Vec<(u32, u32)> = vec![];

    for i in 2.. {
        // Update all the prime marks in sieve for the current number
        // to the next multiple. If there is no mark at all, then the
        // `is_prime` is true.

        let mut is_prime = true;
        for entry in sieve.iter_mut() {
            let (p, ref mut pm) = *entry;

            // OPTIMIZATION: p*p is the first multiple we store of
            // each prime. If p*p > i, then same will be true for all
            // larger primes, so there is no point in checking
            // anymore.
            if p * p > i {
                break;
            }

            if *pm == i {
                is_prime = false;
                *pm = *pm + p;
            }
        }

        // If the number is prime, add a mark for it to the sieve!
        if is_prime {
            // Add (i, i*i) to the prime marks.
            sieve.push((
                i,
                if let Some(square) = i.checked_mul(i) {
                    square
                } else {
                    // i*i overflowed. We won't ever reach this
                    // multiple anyway. So just set it to any past
                    // composite number.
                    4
                },
            ));

            // Finish when we have had enough primes.
            if sieve.len() > n as usize {
                break;
            }
        }
    }

    let (p, _) = sieve.last().cloned().unwrap();
    p
}