Add tests and fix variables.

1 files changed, 7 insertions, 6 deletions

diff --git a/cps.sml b/cps.sml
index 56780a9..adb29c0 100644
--- a/cps.sml
+++ b/cps.sml
@@ -61,17 +61,18 @@ struct
 
   fun hoist (expr : Syntax.cexp) : Syntax.cexp =
     let
-      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 (fs @ acc)
-        | funs (Syntax.CPrimop (_, _, _, ks)) acc = foldl (fn (x, acc) => funs x acc) acc ks
-
       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 _) = 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, []))