coroutine.sml

signature COROUTINE =
sig
  type ('a, 'b, 'r) t

  val await : ('a, 'b, 'a) t
  val yield : 'b -> ('a, 'b, unit) t

  (* Monadic interface and stuffs *)
  val map    : ('c -> 'd) -> ('a, 'b, 'c) t -> ('a, 'b, 'd) t
  val return : 'c -> ('a, 'b, 'c) t
  val >>=    : ('a, 'b, 'c) t * ('c -> ('a, 'b, 'd) t) -> ('a, 'b, 'd) t
  val >>     : ('a, 'b, 'c) t * ('a, 'b, 'd) t -> ('a, 'b, 'd) t

  (* Compose to coroutines. [a --> b] runs b until b needs a value
   * then we run a to a [yield]. [a <-- f] is the reverse.
   *)
  val --> : ('a, 'b, 'd) t * ('b, 'c, 'd) t -> ('a, 'c, 'd) t
  val <-- : ('a, 'b, 'd) t * ('b -> ('b, 'c, 'd) t) -> ('a, 'c, 'd) t

  val run : ('a, 'b, 'r) t -> 'r option
end

structure Coroutine : COROUTINE =
struct
  datatype ('a, 'b, 'r) t
    = PURE of 'r
    | YIELD of 'b * (unit -> ('a, 'b, 'r) t)
    | AWAIT of 'a -> ('a, 'b, 'r) t

  val await = AWAIT PURE
  val return = PURE
  fun yield a = YIELD (a, fn () => return ())

  fun map f c =
    case c of
        PURE r => PURE (f r)
      | YIELD (b, c) => YIELD (b, fn () => map f (c ()))
      | AWAIT c => AWAIT (fn a => map f (c a))

  infix >>= >>
  fun a >>= f =
    case a of
        PURE r => f r
      | YIELD (b, c) => YIELD (b, fn () => c () >>= f)
      | AWAIT c => AWAIT (fn a => c a >>= f)
  fun a >> b = a >>= (fn _ => b)

  infixr -->
  infix <--
  fun l --> r =
    case r of
        PURE r => PURE r
      | YIELD (b, c) => YIELD (b, fn () => l --> c ())
      | AWAIT f => l <-- f
  and l <-- r =
    case l of
        PURE r => PURE r
      | YIELD (b, c) => c () --> r b
      | AWAIT f => AWAIT (fn a => f a <-- r)

  fun run c =
    case c of
        PURE r => SOME r
      | _ => NONE
end

open Coroutine
infix >> >>= <--
infixr -->

fun yieldList [] = yield NONE >> return ()
  | yieldList (x :: xs) = yield (SOME x) >> yieldList xs

fun sum xs =
  let
    val producer = yieldList xs >> return 0
    fun consumer acc =
      await >>= (fn SOME i => consumer (acc + i)
                  | NONE => return acc)
  in
    Option.valOf (run (producer --> consumer 0))
  end