(* stream.sml *)

(* implementation of infinite version of Sieve of Eratosthenes *)

(* because the recursion in the functions from, filter and sieve is
   not well-founded, we have to prove them correct using a
   non-standard technique; see the comment on sieve as an example *)

structure Sieve :> SIEVE =
struct

datatype 'a stream' = Cons of 'a * 'a stream

withtype 'a stream = 'a stream' Susp.susp

(* val from : int -> int stream

   from n returns the stream of all integers that are >= n, listed in
   ascending order *)

fun from n = Susp.delay(fn () => from' n)

(* val from' : int -> int stream' *)

and from' n = Cons(n, from(n + 1))

fun split(s, 0) = (nil, s)
  | split(s, n) =
      let val Cons(x, t) = Susp.force s
          val (ys, u)    = split(t, n - 1)
      in (x :: ys, u) end

(* if infinitely many elements of s satisfy p, then filter p s filters
   out the elements of s that don't satisfy p, maintaining the order
   of the remaining elements *)

fun filter p =
      let (* val filt : 'a stream  -> 'a stream *)

          fun filt s = Susp.delay(fn () => filt'(Susp.force s))

          (* val filt' : 'a stream' -> 'a stream' *)

          and filt'(Cons(x, s)) =
                if p x
                then Cons(x, filt s)
                else filt'(Susp.force s) (* optimization of
                                            Susp.force(filt s) *)
      in filt end

(* divides : int -> int -> bool

   divides n m tests whether n divides m, i.e., whether m is divisible
   by n *)

fun divides (n : int) m = m mod n = 0

(* val sieve  : int stream  -> int stream

   if n is the first element of s, and the elements of s are all the
   integers (in ascending order) that are >= n and are not divisible
   by any integers that are >= 2 and < n, then sieve s returns the
   stream of all primes that are >= n, listed in ascending order

   to prove that sieve s is correct, we use mathematical induction
   to show that, for all m >= 0, split(sieve s, m) = (xs, t),
   where

     xs is the first m primes (in ascending order) that are >= n, and

     t is the stream whose first element, l, is the (m+1)st prime, and
     whose elements are all the integers (in ascending order) that are
     >= l and are not divisible by any integers that are >= 2 and < l *)

fun sieve s = Susp.delay(fn () => sieve'(Susp.force s))

(* val sieve' : int stream' -> int stream' *)

and sieve'(Cons(n, s)) = Cons(n, sieve(filter (not o divides n) s))

val primes = sieve(from 2)

end;