Add tests and fix variables.

1 files changed, 5 insertions, 5 deletions

diff --git a/codegen.sml b/codegen.sml
index 84f3ad5..befedac 100644
--- a/codegen.sml
+++ b/codegen.sml
@@ -51,9 +51,9 @@ struct
       shuffle' inputMap outputMap @ constants
     end
 
-  fun buildVarMap (expr : Syntax.cexp) : Syntax.var VarMap.map =
+  fun buildVarMap (first : int) (expr : Syntax.cexp) : Syntax.var VarMap.map =
     let
-      val next = ref 0
+      val next = ref first
       fun insert v m =
         let val this = !next
         in
@@ -78,7 +78,7 @@ struct
                           (ListPair.zip
                             (List.tabulate (length args, fn x => x + 1),
                              args))
-                      val bodyVars = buildVarMap body
+                      val bodyVars = buildVarMap (length args + 1) body
                     in VarMap.union (VarMap.union acc argsVars) bodyVars
                     end)
                   VarMap.empty
@@ -106,7 +106,7 @@ struct
 
   fun toASM (expr : Syntax.cexp) : Syntax.opcode list =
     let
-      val varMap = buildVarMap expr
+      val varMap = buildVarMap 0 expr
       fun translate v =
         case VarMap.lookup v varMap of
           NONE => raise Fail ("unable to translate var " ^ Int.toString v)
@@ -144,7 +144,7 @@ struct
             in
               bodyASM @ funcsASM
             end
-        | Syntax.CPrimop (Syntax.PExit, [arg], _, _)=> [Syntax.OExit arg]
+        | Syntax.CPrimop (Syntax.PExit, [arg], _, _)=> [Syntax.OExit (translateVal arg)]
         | _ => raise Fail ("malformed CPS:\n" ^ Syntax.cexpToString expr)
     in go expr
     end