module S = Top

type consts_env = int Atom.Map.t

type reduced_value = IntRed of int | ValRed of S.value
let rec reduce_val env (v: S.value) = match v with
| S.VVar var ->
        (try IntRed (Atom.Map.find var env)
        with Not_found -> ValRed (S.VVar var))
| S.VLit x -> IntRed x
| S.VBinOp (left, op, right) ->
        let op_func = function
            | S.OpAdd -> (+)
            | S.OpSub -> (-)
            | S.OpMul -> ( * )
            | S.OpDiv -> (/)
        in
        let apply op l r = match (l, r) with
        | IntRed l, IntRed r -> IntRed ((op_func op) l r)
        | IntRed l, ValRed r -> ValRed (S.VBinOp (S.VLit l, op, r))
        | ValRed l, IntRed r -> ValRed (S.VBinOp (l, op, S.VLit r))
        | ValRed l, ValRed r -> ValRed (S.VBinOp (l, op, r))
        in

        let lVal = reduce_val env left
        and rVal = reduce_val env right in
        apply op lVal rVal

let do_value env value = match reduce_val env value with
| IntRed x -> S.VLit x
| ValRed exp -> exp

let do_block env (S.Con(tag, args)) =
    S.Con (tag, List.map (do_value env) args)

let rec do_term (env: consts_env) (t: S.term) = match t with
| S.Exit -> S.Exit
| S.TailCall (func, args) ->
        S.TailCall (
            func,
            List.map (do_value env) args)
| S.Print (value, next) ->
        S.Print (
            do_value env value,
            do_term env next)
| S.LetVal (var, value, next) ->
        (match do_value env value with
        | S.VLit x -> do_term (Atom.Map.add var x env) next
        | value -> S.LetVal (var, value, do_term env next))
| S.LetBlo (var, block, next) ->
        S.LetBlo (var, do_block env block, do_term env next)
| S.IfZero (cond, tIf, tElse) ->
        S.IfZero (do_value env cond, do_term env tIf, do_term env tElse)
| S.Swi (tag, branches) ->
        S.Swi(tag, List.map (do_branch env) branches)

and do_branch env (S.Branch (tag, args, body)) =
    S.Branch (tag, args, do_term env body)

let do_func (S.Fun (name, args, body)) =
    S.Fun (name, args, do_term Atom.Map.empty body)

let constant_propagation (S.Prog (funcs, body)) =
    S.Prog (List.map do_func funcs, do_term Atom.Map.empty body)