(**
 * 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 Types
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, _), ann) -> BoolConst  (not value, ann)
    | (Neg, IntConst   (value, _), ann) -> IntConst   (-value, ann)
    | (Neg, FloatConst (value, _), ann) -> FloatConst (-.value, ann)
    | (op, opnd, ann) -> Monop (op, opnd, ann)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    (* Constant assignments are added to constants table *)
    | Assign (name, None, value, ann) 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, ann)

    | VarLet (dec, None, value, ann) when is_const_name (nameof dec) ->
        let value = propagate value in
        if is_const value then (
            Hashtbl.add consts (nameof dec) value;
            DummyNode
        ) else
            VarLet (dec, None, value, ann)

    (* Variables that are in the constant table are replaced with their constant
     * value *)
    | Var (name, None, ann) when Hashtbl.mem consts name ->
        Hashtbl.find consts name
    | VarUse (dec, None, ann) when Hashtbl.mem consts (nameof dec) ->
        Hashtbl.find consts (nameof dec)
    | Dim (name, ann) when Hashtbl.mem consts name ->
        Hashtbl.find consts name

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

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

    | Cond (cond, texp, fexp, ann) ->
        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, ann)
        )

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

    | _ -> transform_children propagate node

let rec prune_vardecs consts = function
    | VarDec (ctype, name, init, ann) when Hashtbl.mem consts name -> DummyNode
    | node -> transform_children (prune_vardecs consts) node

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