(**
 * The compiler sometimes generates variables of the form foo$1, to make sure
 * that expressions are only executed once. In many cases, this leads to
 * over-complex constructions, for example when converting for-loops to
 * while-loops. We use the knowledge of these variables being constant by
 * propagation the constant values to their occurrences, and then apply
 * arithmetic simplification to operators to reduce the size and complexity of
 * the generated code. Note that this can only be applied to constants. For
 * variables in general, some form of liveness analysis would be required (e.g.
 * Static Single Assignment form). Expressions can only be propagated when they
 * have no side effects, i.e. when they do not contain function calls.
 *)
open Ast
open Util

let is_const_name name =
    Str.string_match (Str.regexp "[^\\$]+\\$\\$[0-9]+") name 0

let is_const = function
    | BoolConst _ | IntConst _ | FloatConst _ -> true
    | _ -> false

let eval_monop = function
    | (Not, BoolConst  (value, _), loc) -> BoolConst  (not value, loc)
    | (Neg, IntConst   (value, _), loc) -> IntConst   (-value, loc)
    | (Neg, FloatConst (value, _), loc) -> FloatConst (-.value, loc)
    | (op, opnd, loc) -> Monop (op, opnd, loc)

let eval_binop = function
    (* Arithmetic *)
    | (Add, IntConst (left, _), IntConst (right, _), loc) ->
        IntConst (left + right, loc)
    | (Add, FloatConst (left, _), FloatConst (right, _), loc) ->
        FloatConst (left +. right, loc)

    | (Sub, IntConst (left, _), IntConst (right, _), loc) ->
        IntConst (left - right, loc)
    | (Sub, FloatConst (left, _), FloatConst (right, _), loc) ->
        FloatConst (left -. right, loc)

    | (Mul, IntConst (left, _), IntConst (right, _), loc) ->
        IntConst (left * right, loc)
    | (Mul, FloatConst (left, _), FloatConst (right, _), loc) ->
        FloatConst (left *. right, loc)

    | (Div, IntConst (left, _), IntConst (right, _), loc) ->
        IntConst (left / right, loc)
    | (Div, FloatConst (left, _), FloatConst (right, _), loc) ->
        FloatConst (left /. right, loc)

    | (Mod, IntConst (left, _), IntConst (right, _), loc) ->
        IntConst (left mod right, loc)

    (* Relational *)
    | (Eq, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left = right, loc)
    | (Eq, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left = right, loc)

    | (Ne, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left != right, loc)
    | (Ne, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left != right, loc)

    | (Gt, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left > right, loc)
    | (Gt, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left > right, loc)

    | (Lt, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left < right, loc)
    | (Lt, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left < right, loc)

    | (Ge, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left >= right, loc)
    | (Ge, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left >= right, loc)

    | (Le, IntConst (left, _), IntConst (right, _), loc) ->
        BoolConst (left <= right, loc)
    | (Le, FloatConst (left, _), FloatConst (right, _), loc) ->
        BoolConst (left <= right, loc)

    (* Logical *)
    | (And, BoolConst (left, _), BoolConst (right, _), loc) ->
        BoolConst (left && right, loc)
    | (Or, BoolConst (left, _), BoolConst (right, _), loc) ->
        BoolConst (left || right, loc)

    | (op, left, right, loc) -> Binop (op, left, right, loc)

let rec propagate consts node =
    let propagate = propagate consts in
    match node with

    (* Constant assignments are added to constants table *)
    | Assign (name, None, value, loc) when is_const_name name ->
        let value = propagate value in
        if is_const value then (
            Hashtbl.add consts name value;
            DummyNode
        ) else
            Assign (name, None, value, loc)

    (* Variables that are in the constant table are replaced with their constant
     * value *)
    | Var (name, loc) when Hashtbl.mem consts name ->
        Hashtbl.find consts name

    (* Apply arithmetic simplification to constant operands *)
    | Monop (op, opnd, loc) ->
        let opnd = propagate opnd in
        if is_const opnd
            then eval_monop (op, opnd, loc)
            else Monop (op, opnd, loc)

    | Binop (op, left, right, loc) ->
        let left = propagate left in
        let right = propagate right in
        if is_const left && is_const right
            then eval_binop (op, left, right, loc)
            else Binop (op, left, right, loc)

    | Cond (cond, texp, fexp, loc) ->
        let cond = propagate cond in
        let texp = propagate texp in
        let fexp = propagate fexp in
        (match cond with
        | BoolConst (value, _) -> if value then texp else fexp
        | _ -> Cond (cond, texp, fexp, loc)
        )

    | TypeCast (ctype, value, loc) ->
        let value = propagate value in
        (match (ctype, value) with
        | (Bool,  BoolConst  (value, _)) -> BoolConst (value, loc)
        | (Bool,  IntConst   (value, _)) -> BoolConst (value != 1, loc)
        | (Bool,  FloatConst (value, _)) -> BoolConst (value != 1.0, loc)
        | (Int,   BoolConst  (value, _)) -> IntConst ((if value then 1 else 0), loc)
        | (Int,   IntConst   (value, _)) -> IntConst (value, loc)
        | (Int,   FloatConst (value, _)) -> IntConst (int_of_float value, loc)
        | (Float, BoolConst  (value, _)) -> FloatConst ((if value then 1. else 0.), loc)
        | (Float, IntConst   (value, _)) -> FloatConst (float_of_int value, loc)
        | (Float, FloatConst (value, _)) -> FloatConst (value, loc)
        | _ -> TypeCast (ctype, value, loc)
        )

    | _ -> transform_children propagate node

let rec prune_vardecs consts = function
    | VarDec (ctype, name, init, loc) when Hashtbl.mem consts name -> Block []
    | node -> transform_children (prune_vardecs consts) node

let rec phase input =
    prerr_endline "- Constant propagation";
    match input with
    | Ast node ->
        let consts = (Hashtbl.create 32) in
        let node = propagate consts node in
        Ast (prune_vardecs consts node)
    | _ -> raise (InvalidInput "constant propagation")