diff options
Diffstat (limited to 'CPS.sml')
| -rw-r--r-- | CPS.sml | 271 |
1 files changed, 271 insertions, 0 deletions
@@ -0,0 +1,271 @@ +structure CPS = +struct + fun toCPS (e : Syntax.lexp) (cont : Syntax.value -> Syntax.cexp) : Syntax.cexp = + case e of + Syntax.LVar v => cont (Syntax.VVar v) + | Syntax.LFn (v, expr) => + let + val fnName = GenSym.new () + val k = GenSym.new () + in + Syntax.CFix + ([(fnName, [v, k], toCPS expr (fn ret => Syntax.CApp (Syntax.VVar k, [ret])))], + cont (Syntax.VVar fnName)) + end + | Syntax.LFix (decls, body) => + Syntax.CFix + ( map + (fn (name, arg, expr) => + let val w = GenSym.new () + in + ( name + , [arg, w] + , toCPS expr (fn z => Syntax.CApp (Syntax.VVar w, [z])) + ) + end) + decls + , toCPS body cont + ) + | Syntax.LApp (Syntax.LPrim primop, Syntax.LRecord args) => + let + val temp = GenSym.new () + fun go [] acc = Syntax.CPrimop (primop, rev acc, [temp], [cont (Syntax.VVar temp)]) + | go (arg :: args) acc = toCPS arg (fn arg' => go args (arg' :: acc)) + in go args [] + end + | Syntax.LApp (Syntax.LPrim primop, arg) => toCPS (Syntax.LApp (Syntax.LPrim primop, Syntax.LRecord [arg])) cont + | Syntax.LApp (f, x) => + let + val addr = GenSym.new () + val arg = GenSym.new () + in Syntax.CFix + ([(addr, [arg], cont (Syntax.VVar arg))], + toCPS f (fn f' => + toCPS x (fn x' => + Syntax.CApp (f', [x', Syntax.VVar addr])))) + end + | Syntax.LInt i => cont (Syntax.VInt i) + | Syntax.LString s => + let val temp = GenSym.new () + in + Syntax.CRecord + ( [(map (fn c => (Syntax.VInt (Char.ord c), [])) (String.explode s), temp)] + , cont (Syntax.VVar temp) + ) + end + | Syntax.LRecord [] => cont (Syntax.VInt 0) + | Syntax.LSelect (i, expr) => + let val temp = GenSym.new () + in toCPS expr (fn x => Syntax.CSelect (i, x, temp, cont (Syntax.VVar temp))) + end + | Syntax.LRecord exprs => + let + fun go [] vars = + let val temp = GenSym.new () + in Syntax.CRecord ([(map (fn v => (v, [])) (rev vars), temp)], cont (Syntax.VVar temp)) + end + | go (expr :: exprs) vars = + toCPS expr (fn v => go exprs (v :: vars)) + in go exprs [] + end + | Syntax.LSwitch (expr, arms, otherwise) => + let + val sortedArms = Sort.sort (fn ((x, _), (y, _)) => Int.compare (x, y)) arms + val contAddr = GenSym.new () + val otherwiseAddr = GenSym.new () + val arg = GenSym.new () + fun go _ [] cont = Syntax.CApp (Syntax.VVar otherwiseAddr, []) + | go v [(x, arm)] cont = + (case otherwise of + NONE => toCPS arm cont + | SOME _ => + let val b = GenSym.new () + in + Syntax.CPrimop + ( Syntax.PEq + , [v, Syntax.VInt x] + , [b] + , [ Syntax.CPrimop + ( Syntax.PIf + , [Syntax.VVar b] + , [] + , [toCPS arm cont, Syntax.CApp (Syntax.VVar otherwiseAddr, [])] + ) + ] + ) + end) + | go v arms cont = + let + val b = GenSym.new () + val h = length arms div 2 + val half1 = List.take (arms, h) + val half2 = List.drop (arms, h) + val (x, _) = hd half2 + in + Syntax.CPrimop + ( Syntax.PLess + , [v, Syntax.VInt x] + , [b] + , [ Syntax.CPrimop + ( Syntax.PIf + , [Syntax.VVar b] + , [] + , [go v half1 cont, go v half2 cont] + ) + ] + ) + end + fun contFunc x = Syntax.CApp (Syntax.VVar contAddr, [x]) + val fixFuncs = [(contAddr, [arg], cont (Syntax.VVar arg))] + val fixFuncs = + case otherwise of + NONE => fixFuncs + | SOME otherwise => (otherwiseAddr, [], toCPS otherwise contFunc) :: fixFuncs + in + Syntax.CFix + ( fixFuncs + , toCPS expr (fn v => go v sortedArms contFunc) + ) + end + | _ => raise Fail ("malformed expression " ^ Syntax.lexpToString e) + + fun hoist (expr : Syntax.cexp) : Syntax.cexp = + let + fun exprs (Syntax.CRecord (records, k)) = Syntax.CRecord (records, exprs k) + | exprs (Syntax.CSelect (i, arg, res, k)) = Syntax.CSelect (i, arg, res, exprs k) + | exprs (expr as Syntax.CApp _) = expr + | exprs (Syntax.CFix (_, body)) = exprs body + | exprs (Syntax.CPrimop (p, args, res, ks)) = Syntax.CPrimop (p, args, res, (map exprs ks)) + + fun funs (Syntax.CRecord (_, k)) acc = funs k acc + | funs (Syntax.CSelect (_, _, _, k)) acc = funs k acc + | funs (Syntax.CApp _) acc = acc + | funs (Syntax.CFix (fs, body)) acc = + funs body (foldl (fn ((fName, fVars, fBody), acc) => (fName, fVars, exprs fBody) :: funs fBody acc) acc fs) + | funs (Syntax.CPrimop (_, _, _, ks)) acc = foldl (fn (x, acc) => funs x acc) acc ks + val entryPoint = GenSym.new () + in + Syntax.CFix ((entryPoint, [], exprs expr) :: funs expr [], Syntax.CApp (Syntax.VLabel entryPoint, [])) + end + + structure VarMap = Map (type k = Syntax.var + val cmp = Int.compare) + + fun varSet (l : Syntax.var list) : unit VarMap.map = VarMap.fromList (map (fn x => (x, ())) l) + + fun freeVars (expr : Syntax.cexp) : unit VarMap.map = + case expr of + Syntax.CRecord (records, k) => + let + val kFreeVars = freeVars k + val argFreeVars = + foldl + (fn ((args, _), acc) => + VarMap.union acc (varSet (List.mapPartial (fn (Syntax.VVar v, _) => SOME v | _ => NONE) args))) + VarMap.empty + records + val results = + foldl + (fn ((_, res), acc) => + VarMap.insert res () acc) + VarMap.empty + records + in + VarMap.difference (VarMap.union kFreeVars argFreeVars) results + end + | Syntax.CSelect (_, arg, res, k) => + let + val argFreeVars = + case arg of + Syntax.VVar v => varSet [v] + | _ => VarMap.empty + val kFreeVars = VarMap.delete res (freeVars k) + in + VarMap.union argFreeVars kFreeVars + end + | Syntax.CApp (func, args) => + let + val funcFreeVars = + case func of + Syntax.VVar v => varSet [v] + | _ => VarMap.empty + val argFreeVars = varSet (List.mapPartial (fn Syntax.VVar v => SOME v | _ => NONE) args) + in + VarMap.union funcFreeVars argFreeVars + end + | Syntax.CFix (funs, body) => + let + val names = varSet (map (fn (name, _, _) => name) funs) + val funsFreeVars = map (fn (name, args, fixBody) => VarMap.difference (freeVars fixBody) (varSet args)) funs + val bodyFreeVars = freeVars body + in + foldl (fn (x, acc) => VarMap.union acc (VarMap.difference x names)) VarMap.empty (bodyFreeVars :: funsFreeVars) + end + | Syntax.CPrimop (_, args, res, ks) => + let + val argFreeVars = varSet (List.mapPartial (fn Syntax.VVar v => SOME v | _ => NONE) args) + val boundVars = varSet res + val kFreeVars = foldl (fn (x, acc) => VarMap.union acc x) VarMap.empty (map (fn k => VarMap.difference (freeVars k) boundVars) ks) + in + VarMap.union argFreeVars kFreeVars + end + + fun freeVarsClosure (name, args, body) = map (fn (x, _) => x) (VarMap.toList (VarMap.difference (freeVars body) (varSet args))) + + fun enumerate l = ListPair.zip (List.tabulate (length l, (fn x => x)), l) + + fun convertExpr varMap expr = + let + fun translate var = getOpt (VarMap.lookup var varMap, var) + fun translateValue (Syntax.VVar v) = Syntax.VVar (translate v) + | translateValue v = v + in + case expr of + Syntax.CRecord (records, k) => + Syntax.CRecord (map (fn (args, res) => (map (fn (v, p) => (translateValue v, p)) args, res)) records, convertExpr varMap k) + | Syntax.CSelect (i, arg, res, k) => + Syntax.CSelect (i, translateValue arg, res, convertExpr varMap k) + | Syntax.CApp (func, args) => + let + val temp = GenSym.new () + val f = translateValue func + in Syntax.CSelect (0, f, temp, + Syntax.CApp (Syntax.VVar temp, f :: map translateValue args)) + end + | Syntax.CFix (funcs, body) => + let + val convertedFuncs = + map + (fn this as (name, args, body) => + let + val funcFreeVars = freeVarsClosure this + val varMap' = VarMap.union varMap (VarMap.fromList (map (fn v => (v, GenSym.new ())) funcFreeVars)) + val closure = GenSym.new () + val newBody = + foldl + (fn ((i, x), acc) => Syntax.CSelect (i + 1, Syntax.VVar closure, valOf (VarMap.lookup x varMap'), acc)) + (convertExpr varMap' body) + (enumerate funcFreeVars) + in + (GenSym.new (), closure :: args, newBody) + end) + funcs + val newBody = + Syntax.CRecord + ( map + (fn (old as (oldName, args, body), (newName, _, _)) => + let val funcFreeVars = freeVarsClosure old + in ((Syntax.VLabel newName, []) :: map (fn v => (Syntax.VVar (translate v), [])) funcFreeVars, oldName) + end) + (ListPair.zip (funcs, convertedFuncs)) + , convertExpr varMap body + ) + in + Syntax.CFix (convertedFuncs, newBody) + end + | Syntax.CPrimop (p, args, res, ks) => + Syntax.CPrimop (p, map translateValue args, res, map (convertExpr varMap) ks) + end + + fun convertClosures (expr : Syntax.cexp) : Syntax.cexp = hoist (convertExpr VarMap.empty expr) +end |
