Remove var-var bindings after defun

src/Main.ml

@@ -95,6 +95,7 @@ let process filename =
   |> dump "Tail" Tail.show_term false
   |> dump "PrettyTail" PrettyTail.show true
   |> Defun.defun_term
+  |> VarVarBind.clean_var_var_bind
   |> dump "Top" Top.show_program false
   |> dump "PrettyTop" PrettyTop.show true
   |> Finish.finish_program

src/VarVarBind.ml

@@ -0,0 +1,63 @@
+module S = Top
+module Env = Atom.Map
+
+let empty_env = Env.empty
+
+let clean_var map_env var =
+    try Env.find var map_env
+    with Not_found -> var
+
+let rec clean_value map_env = function
+| S.VVar var -> S.vvar @@ clean_var map_env var
+| S.VLit x -> S.VLit x
+| S.VBinOp (v1, op, v2) ->
+        S.VBinOp (
+            clean_value map_env v1,
+            op,
+            clean_value map_env v2)
+
+let rec clean_term map_env = function
+| S.Exit -> S.Exit
+| S.TailCall (func, args) ->
+        S.TailCall (
+            clean_var map_env func,
+            List.map (clean_value map_env) args)
+| S.Print (value, next) ->
+        S.Print (
+            clean_value map_env value,
+            clean_term map_env next)
+| S.LetVal (name, S.VVar subVar, next) ->
+        clean_term (Env.add name subVar map_env) next
+| S.LetVal (name, value, next) ->
+        S.LetVal (name, value, clean_term map_env next)
+| S.LetBlo (name, S.Con(tag, args), next) ->
+        S.LetBlo(
+            name,
+            S.Con(tag, List.map (clean_value map_env) args),
+            clean_term map_env next)
+| S.IfZero (cond, tIf, tElse) ->
+        S.IfZero (
+            clean_value map_env cond,
+            clean_term map_env tIf,
+            clean_term map_env tElse)
+| S.Swi (tag, branches) ->
+        S.Swi (
+            tag,
+            List.map (clean_branch map_env ) branches)
+
+and clean_branch map_env (S.Branch(tag, args, body)) =
+    S.Branch(
+        tag,
+        args,
+        clean_term map_env body)
+
+let clean_function (S.Fun(name, args, body)) =
+    S.Fun(
+        name,
+        args,
+        clean_term empty_env body)
+
+let clean_var_var_bind (S.Prog(funcs, body)) =
+    S.Prog(
+        List.map clean_function funcs,
+        clean_term empty_env body)

src/VarVarBind.mli

@@ -0,0 +1,2 @@
+(** Removes var-var bindings, eg. let x = y in e, where x and y are vars *)
+val clean_var_var_bind: Top.program -> Top.program