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 (clean_var map_env subVar) map_env) next
| S.LetVal (name, value, next) ->
        S.LetVal (name, clean_value map_env 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,
        List.map (clean_var map_env) 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)