structure Lazy =
struct
datatype 'a t_ = T of unit -> 'a
               | V of 'a
               | E of exn
withtype 'a t = 'a t_ ref

fun lazy t = ref (T t)
fun eager v = ref (V v)
fun force p =
    case !p of
      V v => v
    | E e => raise e
    | T t =>
      let
        val v = t ()
      in
        p := V v ; v
      end
      handle e => (p := E e ; raise e)
fun delay t = lazy (fn _ => force (t ()))

fun thunk p () = force p
fun memoise t = thunk (lazy t)

exception Undefined
val undefined = lazy (fn _ => raise Undefined)
end

type 'a lazy = 'a Lazy.t
val F = Lazy.force
val E = Lazy.eager
val L = Lazy.lazy
val D = Lazy.delay
val undefined = Lazy.undefined

infixr $
fun f $ a = f a
fun fst (x, _) = x
fun snd (_, x) = x

datatype 'a stream = Stream of ('a * 'a stream lazy)

fun take 0 _ = nil
  | take n s =
    let
      val Stream (x, s') = F s
    in
      x :: take (n - 1) s'
    end

fun map f s =
    L (fn _ =>
          let
            val Stream (x, s') = F s
          in
            Stream (f x, map f s')
          end
      )

fun scan f b s =
    L (fn _ =>
          let
            val Stream (x, s') = F s
            val y = f (x, b)
          in
            Stream (y, scan f y s')
          end
      )

fun zip s t =
    L (fn _ =>
          let
            val Stream (x, s') = F s
            val Stream (y, t') = F t
          in
            Stream ((x, y), zip s' t')
          end
      )

val nat =
    let
      fun loop n = L (fn _ => Stream (n, loop $ n + 1))
    in
      loop 1
    end
fun repeat n = L (fn _ => Stream (n, repeat n))
fun cons x s = L (fn _ => Stream (x, s))

type ('s, 'a) rstate = 's lazy -> 'a lazy * 's lazy
infix >>=

val return : ('a lazy -> ('s, 'a) rstate) =
 fn x => fn s => (x, s)

val op>>= : ('s, 'a) rstate * ('a lazy -> ('s, 'b) rstate) -> ('s, 'b) rstate =
 fn (m, k) =>
 fn s =>
    let
      val r = ref undefined
      val s' = L (fn _ => F $ !r)
      val (x, s'') = m s'
      val t =
          L (fn _ =>
                let
                  val (y, s') = k x s
                in
                  r := s' ; y
                end)
      val s'' =
          D (fn _ => (F t ; s''))
      val y =
          D (fn _ => F t)
    in
      (y, s'')
    end

val get : ('s, 's) rstate =
 fn s => (s, s)

val put : 's lazy -> ('s, unit) rstate =
 fn s => fn _ => (E (), s)

val modify =
 fn f => fn s => (E (), f s)

val yield : 'a -> ('a stream, unit) rstate =
 fn x =>
 fn s =>
    (E (), cons x s)

fun fiba x y =
    do yield x
     ; fiba y (x + y)
    end

fun cumsum s =
    let
      fun loop n s =
          cons n $
               D (fn _ =>
                     let
                       val Stream (x, s') = F s
                     in
                       loop (n + x) s'
                     end
                 )
    in
      loop 0 s
    end

val fibb =
    do fibs <- get
     ; modify cumsum
     ; put $ cons 1 fibs
     ; return fibs
    end

val fact =
    do f <- get
     ; modify (map op* o zip nat)
     ; put $ cons 1 f
     ; return f
    end

; take 10 $ snd $ fiba 0 1 undefined ;
; take 10 $ snd $ fibb undefined ;