cps.sml

signature NAME = sig
  type t
  val empty: string -> t
  val fresh: t -> string  
end;

structure Name = struct
  datatype t = T of {name: string, id: int ref};  
  fun empty name = T {name=name, id=ref 0};  
  fun fresh (T {name, id}) = 
    let 
      val x = !id; 
      val () = id := x + 1; 
    in name ^ (Int.toString x) end
end :> NAME;
  
structure ML = struct 
  datatype t = 
    Var of string 
    | App of t * t
    | Lam of string * t
    | Pair of t * t
    | Proj of t * int 
    | Unit
    | Int of int
    | Let of string * t * t
    | LetMany of (string * t) list * t
  
  val rec pp = 
   fn Var s => print s
    | App (e1, e2) => (pp e1; print " "; pp e2)
    | Lam (s, e) => (print ("fn " ^ s ^ " => "); pp e)
    | Pair (e1, e2) => (print "("; pp e1; print ","; pp e2; print ")")
    | Proj (e1, idx) => (print "#"; print (Int.toString idx); print " "; pp e1)
    | Unit => print "()"
    | Int i => print (Int.toString i)
    | Let (b, expr, body) => 
      (print ("let " ^ b ^ " = "); pp expr; print " in "; pp body; print " end")
end;

structure CPS = struct  
  datatype cps
    (* letval x = val in body *)
    = LetV of string * cval * cps
    (* letproj x = #1 x in K *)
    | LetP of string * string * int * cps
    (* letkont k x = kexpr in body *)
    | LetC of string * string * cps * cps
    (* continue with arg *)
    | KApp of {kont: string, arg: string}
    (* call with arg *)
    | FApp of {func: string, kont: string, arg: string}
  and cval
    = Unit
    | Var of string
    | Int of int
    | Pair of cval * cval
    | CLam of {kont: string, arg: string, body: cps};


  fun mkfresh x = let val n = Name.empty x in fn () => Name.fresh n end  
  val freshk = mkfresh "k"
  val freshx = mkfresh "x"
  
  fun transform (kcontext : string -> cps) =
   fn ML.Var s => kcontext s
    | ML.Unit => let val bind = freshx () in LetV (bind, Unit, kcontext bind) end
    | ML.App (e1, e2) => transform (fn z1 => transform (fn z2 => 
        let 
            val k = freshk () 
            val x = freshx ()
        in 
        LetC (k, x, kcontext x, FApp { func=z1, kont=k, arg=z2}) end) e2) e1
    | ML.Lam (x, e) => let 
          val f = freshx () 
          val k = freshk ()
          val value = CLam { kont=k, arg = x, body = transform (fn z => KApp {kont=k, arg=z}) e}
        in 
          LetV (f, value, kcontext f)
        end 
    | ML.Pair (e1, e2) =>  transform (fn z1 => transform (fn z2 => 
        let 
            val x = freshx ()
        in 
        LetV (x, Pair (Var z1, Var z2), kcontext x) end) e2) e1
    | ML.Proj (e1, idx) =>  transform (fn z => let val x = freshx () in LetP (x, z, idx, kcontext x) end) e1
    | ML.Int i =>  let val x = freshx () in LetV (x, Int i, kcontext x) end
    | ML.Let (x, bind, body) =>  let 
        val k = freshk () 
        val body = transform kcontext body
        val bind = transform (fn z => KApp {kont=k, arg=z}) bind 
        in 
            LetC (k, x, body, bind)
        end
    | ML.LetMany (xs, body) => 
      let 
        val k = freshk()
        val expr = transform kcontext body 
        fun inner ((x, bind), (k, body)) = 
          let 
            val kk = freshk()
            val expr = transform (fn z => KApp {kont=k, arg=z}) bind
          in 
            (kk, LetC (kk, x, body, expr))
          end
      in 
        #2 (List.foldr inner (k, expr) xs)
      end

    val rec pp_cps = 
     fn LetV (bind, value, body) => (print "letv "; print bind; print " = "; pp_val value; print " in "; pp_cps body; print " end")
      | LetP (bind, value, index, body) => 
        (print "letproj "; print bind; print " = #"; print (Int.toString index); 
         print " "; print value; print " in "; pp_cps body; print " end")
      | LetC (kont, arg, expr, body) => (print "letcont "; print kont; print " "; print arg; print " = "; pp_cps expr; print " in "; pp_cps body; print " end")
      | KApp {kont, arg} => (print kont; print " "; print arg)
      | FApp {func, kont, arg} => print (func ^ " " ^ kont ^ " " ^ arg)
    and pp_val = 
     fn Unit => print "()"
      | Var s => print s
      | Int i => print (Int.toString i)
      | Pair (e1, e2) => (print "("; pp_val e1; print ", "; pp_val e2; print ")")
      | CLam {kont, arg, body} => (print ("lam " ^ kont ^ " " ^ arg ^ " => "); pp_cps body)
end


val ex = ML.Lam ("x", ML.App (ML.Var "f", ML.Pair (ML.Var "x", ML.Var "y")))
val ex2 = ML.Let("x", ML.Int 8, ML.Let ("y", ML.Int 10, ML.Pair (ML.Var "x", ML.Var "y")))

val ex3 = ML.LetMany ([
  ("x", ML.Int 8),
  ("y", ML.Int 10)
  (* ("z", ML.Int 7) *)
], ML.Pair (ML.Var "x", ML.Var "y"))

val run = CPS.pp_cps o CPS.transform (fn x => CPS.KApp {kont="halt", arg=x})

val () = run ex