(* Source code to                                                       *)
(*   Edoardo Biagioni, Ken Cline, Peter Lee, Chris Okasaki, Chris Stone *)
(*   Safe-for-Space Threads in Standard ML                              *)
(*   CW'97                                                              *)

(* uses structure Queue from the SML/NJ library *)

signature COROUTINE =
sig
   exception MainThreadCantExit
   exception ChildThreadCantSync 

   val fork : (unit -> unit) -> unit
   val yield: unit -> unit
   val exit : unit -> 'a 

   val sync : unit -> unit
     (* sync () yields until all other threads have completed *)
end 

structure Coroutine : COROUTINE =
struct
   exception MainThreadCantExit
   exception ChildThreadCantSync 

   val callcc = SMLofNJ.callcc
   val throw  = SMLofNJ.throw

   datatype threadType = Main | Child
   type thread = unit cont * threadType 

   val readyQueue : thread Queue.queue = Queue.mkQueue ()
   val syncCont : thread option ref = ref NONE
   val currentThreadType = ref Main 

   fun enqueue thread = Queue.enqueue (readyQueue, thread) 

   fun dispatch () =
      let val (t,typ) = 
             Queue.dequeue readyQueue
             handle Queue.Dequeue =>
                (* syncCont cannot be NONE *)
                case !syncCont of SOME main => main
      in 
          currentThreadType := typ;
          throw t () 
      end 

   fun exit () =
      case !currentThreadType of
         Main  => raise MainThreadCantExit
       | Child => dispatch () 

   fun sync () =
      case !currentThreadType of
         Main => callcc (fn t =>
                   (syncCont := SOME (t, Main);
                    dispatch ()))
       | Child => raise ChildThreadCantSync 

   fun yield () =
      callcc (fn parent =>
        (enqueue (parent, !currentThreadType);
         dispatch ()))

   val threadActivator : (unit -> unit) cont =
      callcc (fn return =>
        (let val f = callcc (fn fc => throw return fc)
         in f () end
         handle _ => print "Uncaught exception.";
         exit ();
         (* raise dummy exception as hint to compiler *)
         raise MainThreadCantExit)) 

   fun fork f =
      callcc (fn parent =>
        (enqueue (parent, !currentThreadType);
         currentThreadType := Child;
         throw threadActivator f)) 
end