(* primes.sml *)

(* implementation of generation of primes *)

structure Primes :> PRIMES =
struct

type state = int list * int option * int list * int

(* invariant on state of the form (ms, p, ls, k): ms @ rev ls is the
   first length ms + length ns primes in ascending order, ms is
   nonempty, p is SOME of the square of the last element of ms, or
   NONE if that square would result in overflow, and k is the last
   element of ms @ rev ls *)

val init = ([2], SOME 4, [], 2) : state

(* val divisible : int * int list -> bool

   if ms is an ascending sequence of primes with no gaps, and, if ms
   <> [], then l is not divisible by any primes < hd ms, and l is >
   the last element of ms, then divisible(l, ms) tests whether l is
   divisible by an element of ms

   it returns true as soon as the next element divides l; otherwise it
   returns false if the square of that next element is > l

   tail recursive, using tail of second argument *)

fun divisible(_, nil)     = false
  | divisible(l, m :: ms) =
      l mod m = 0 orelse
      (m * m < l handle _ => false) andalso divisible(l, ms)

(* val nxt :
         int list * int option * int list * int ->
         int list * int option * int list * int

   if ms @ rev ls are the first length ms + length ls primes, listed
   in ascending order, ms is nonempty and p is SOME of the square of
   the last element of ms, or NONE if that square would result in
   overflow, k is > the last element of ms @ rev ls, and there are no
   primes that are > the last element of ms @ rev ls but < k, then
   nxt(ms, p, ls, k) returns (ms', p', ls', k') where k' is the
   smallest prime that is >= k (and so the smallest prime that is >
   the elements of ms @ rev ls), ms @ rev ls = ms' @ rev ls' are the
   first length ms' + length ls' primes, listed in ascending order,
   ms' is nonempty, p' is SOME of the square of the last element of
   ms', or NONE if that square would result in overflow, k' is > the
   last element of ms' @ rev ls', and there are no primes that are >
   the last element of ms' @ rev ls' but < k'

   tail recursive; distance between fourth argument and smallest
   prime >= fourth argument goes down *)

fun nxt(ms, p, ls, k) =
      let val (ms, p, ls) =
                (* reorganize ms/ls if we need to consider elements
                   of ls (as well as ms) to be sure whether k is prime

                   e.g., when generating the first 5,000,000 primes,
                   the reorganization only happens five times (when k
                   is 5, 10, 50, 2210 or 4870850) *)
                if null ls orelse p = NONE orelse valOf p >= k
                then (ms, p, ls)
                else (ms @ rev ls,
                      SOME(hd ls * hd ls) handle _ => NONE,
                      nil)                      
      in if divisible(k, ms)
         then nxt(ms, p, ls, k + 1)
         else (ms, p, ls, k)
      end

fun next ((ms, p, ls, k) : state) : int * state =
      (k,
       let val (ms, p, ls, k) = nxt(ms, p, ls, k + 1)
           in (ms, p, k :: ls, k) end)

end;