Add lambda functions.

1 files changed, 14 insertions, 4 deletions

diff --git a/cps.sml b/cps.sml
index 948b8c8..6854545 100644
--- a/cps.sml
+++ b/cps.sml
@@ -2,7 +2,17 @@ structure CPS =
 struct
   fun toCPS (e : Syntax.lexp) (cont : Syntax.value -> Syntax.cexp) : Syntax.cexp =
     case e of
-      Syntax.LApp (Syntax.LPrim primop, Syntax.LRecord args) =>
+      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.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)])
@@ -45,8 +55,8 @@ struct
 
       fun exprs (Syntax.CRecord (args, res, k)) = Syntax.CRecord (args, res, exprs k)
         | exprs (Syntax.CSelect (i, arg, res, k)) = Syntax.CSelect (i, arg, res, exprs k)
-        | exprs (expr as Syntax.CApp (func, args)) = expr
-        | exprs (Syntax.CFix (_, body)) = body
+        | 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))
       val entryPoint = Gensym.new ()
     in
@@ -148,7 +158,7 @@ struct
                   let val funcFreeVars = freeVarsClosure old
                   in Syntax.CRecord ((Syntax.VLabel newName, []) :: map (fn v => (Syntax.VVar v, [])) funcFreeVars, oldName, acc)
                   end)
-                body
+                (convertExpr varMap body)
                 (ListPair.zip (funcs, convertedFuncs))
           in
             Syntax.CFix (convertedFuncs, newBody)