(* https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes *)

let debug = ref true ;;

let eratos n =
  (* set the upper bound of numbers to check *)
  let upper = n * n + 1 in

  (* populate a list with numbers that we can record primeness of *)
  let primes = Array.create upper true in
  (* walk the numbers between 2 and the upper bound *)
  for this_num = 2 to upper do
      (* next, just iterate through the sieve *)
      (* yes, this will hit the same numbers bunches of times, but
       * this isn't about efficiency... *)
      for iterator = this_num to upper do
        let multiple = ref 0 in
        let it_sqd = iterator * iterator in
        (* let next_nonprime = iterator * iterator + !multiple * iterator in *)
          (* while next_nonprime < upper do *)
          try
            while it_sqd + !multiple * it_sqd < Array.length primes do

              if !debug then
                let this_prime = it_sqd + !multiple * it_sqd in
                Printf.printf "marking %d non-prime\n" this_prime ;

              (* well, all these numbers are certainly not prime! *)
              primes.( it_sqd + !multiple * it_sqd ) <- false ;
              multiple := !multiple + 1 ;
            done
          with
            (* this shouldn't happen, but if it does ... *)
            Invalid_argument a -> Printf.printf "%s: record (%d * %d)\n" a !multiple it_sqd;
      done
  done ;

  (* bernard on irc points out that 0 and 1 are not primes here *)
  primes.( 0 ) <- false ;
  primes.( 1 ) <- false ;

  (* send off our list to the user *)
  primes ;;

let nth_prime n =
  let myprimes = Stack.create () in
  let primestatus = eratos n in
  for index = 0 to Array.length primestatus - 1 do
    try
    if primestatus.( index ) = true then
      if !debug then
        Printf.printf "selecting %d as prime\n" index ;
      Stack.push index myprimes ;
    with
      (* again, this shouldn't happen, but ... *)
      Invalid_argument a -> Printf.printf "nth_prime: %s: %d\n" a index ;
  done ;
  Stack.pop myprimes ;