(**
 * 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.
 *
 * Constant propagation is merged with some some arithmetic simplification here,
 * specifically targeting optimization oppertunities created bij earlier
 * constant propagation. This is utilized, for example, in array index
 * calculation when array dimensions are constant.
 *)
open Types
open Util

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

let is_const = function Const _ -> true | _ -> false

(* Play-it-safe side effect analysis: only return true for variables and
 * constants, since these are targeted in arithmetic simplification (in
 * particular targeting array indices that can be simplified after array
 * dimension reduction). *)
let no_side_effect = function
  | VarUse _ | Const _ | Var _ -> true
  | _ -> false

(* Constand folding *)
let eval = function
  (* Binop - arithmetic *)
  | Binop (Add, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (IntVal (left + right), ann)
  | Binop (Add, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (FloatVal (left +. right), ann)

  | Binop (Sub, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (IntVal (left - right), ann)
  | Binop (Sub, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (FloatVal (left -. right), ann)

  | Binop (Mul, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (IntVal (left * right), ann)
  | Binop (Mul, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (FloatVal (left *. right), ann)

  | Binop (Div, Const (IntVal left, _), Const (IntVal right, _), ann) when right != 0 ->
    Const (IntVal (left / right), ann)
  | Binop (Div, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (FloatVal (left /. right), ann)

  | Binop (Mod, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (IntVal (left mod right), ann)

  (* Binop - relational *)
  | Binop (Eq, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left = right), ann)
  | Binop (Eq, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left = right), ann)

  | Binop (Ne, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left <> right), ann)
  | Binop (Ne, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left <> right), ann)

  | Binop (Gt, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left > right), ann)
  | Binop (Gt, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left > right), ann)

  | Binop (Lt, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left < right), ann)
  | Binop (Lt, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left < right), ann)

  | Binop (Ge, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left >= right), ann)
  | Binop (Ge, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left >= right), ann)

  | Binop (Le, Const (IntVal left, _), Const (IntVal right, _), ann) ->
    Const (BoolVal (left <= right), ann)
  | Binop (Le, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
    Const (BoolVal (left <= right), ann)

  (* Binop - logical *)
  | Binop (And, Const (BoolVal left, _), Const (BoolVal right, _), ann) ->
    Const (BoolVal (left && right), ann)
  | Binop (Or, Const (BoolVal left, _), Const (BoolVal right, _), ann) ->
    Const (BoolVal (left || right), ann)

  (* Monary operations *)
  | Monop (Not, Const (BoolVal  value, _), ann) -> Const (BoolVal  (not value), ann)
  | Monop (Neg, Const (IntVal   value, _), ann) -> Const (IntVal   (-value), ann)
  | Monop (Neg, Const (FloatVal value, _), ann) -> Const (FloatVal (-.value), ann)

  (* 0 * a --> 0 *)
  | Binop (Mul, Const (IntVal 0, _), other, ann)
  | Binop (Mul, other, Const (IntVal 0, _), ann) when no_side_effect other ->
    Const (IntVal 0, ann)

  (* 0 + a --> a *)
  | Binop (Add, Const (IntVal 0, _), other, _)
  | Binop (Add, other, Const (IntVal 0, _), _) ->
    other

  (* 1 * a --> a *)
  | Binop (Mul, Const (IntVal 1, _), other, _)
  | Binop (Mul, other, Const (IntVal 1, _), _) ->
    other

  (* true|false ? texp : fexp --> texp|fexp*)
  | Cond (Const (BoolVal value, _), texp, fexp, _) ->
    if value then texp else fexp

  | node -> node

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 begin
      Hashtbl.add consts name value;
      DummyNode
    end 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 begin
      Hashtbl.add consts (nameof dec) value;
      DummyNode
    end 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) ->
    eval (Monop (op, propagate opnd, ann))

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

  | Cond (cond, texp, fexp, ann) ->
    eval (Cond (propagate cond, propagate texp, propagate fexp, ann))

  | TypeCast (ctype, value, ann) ->
    let value = propagate value in
    begin match (ctype, value) with
    | (Bool,  Const (BoolVal  value, _)) -> Const (BoolVal value, ann)
    | (Bool,  Const (IntVal   value, _)) -> Const (BoolVal (value != 1), ann)
    | (Bool,  Const (FloatVal value, _)) -> Const (BoolVal (value != 1.0), ann)
    | (Int,   Const (BoolVal  value, _)) -> Const (IntVal (if value then 1 else 0), ann)
    | (Int,   Const (IntVal   value, _)) -> Const (IntVal value, ann)
    | (Int,   Const (FloatVal value, _)) -> Const (IntVal (int_of_float value), ann)
    | (Float, Const (BoolVal  value, _)) -> Const (FloatVal (if value then 1. else 0.), ann)
    | (Float, Const (IntVal   value, _)) -> Const (FloatVal (float_of_int value), ann)
    | (Float, Const (FloatVal value, _)) -> Const (FloatVal value, ann)
    | _ -> TypeCast (ctype, value, ann)
    end

  | _ -> transform_children propagate node

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

let phase = function
  | Ast node ->
    let consts = Hashtbl.create 32 in
    let node = propagate consts node in
    Ast (prune_vardecs consts node)
  | _ -> raise (InvalidInput "constant propagation")