(*
 * Do a number of checks:
 * - A void function must not return a value.
 * - A non-void function must return a value of the correct type.
 * - Array indices must be of type integer.
 * - The number of array indices must match the number of array dimensions.
 * - The type on the right-hand side of an assignment must match the type on
 *   the left-hand side.
 * - The number of arguments used for a function call must match the number of
 *   parameters for that function.
 * - The types of the function arguments must match the types of parameters.
 * - The operands of a unary or binary operation must have valid types.
 * - The predicate expression of an if, while, or do-while statement must be
 *   a boolean.
 * - Only values of a basic type can be type cast.
 *)
open Printf
open Types
open Util
open Stringify

let node_error (node, msg) = FatalError (NodeMsg (node, msg))

(* Stringify a list of types for use in error messages.
 * ctype list -> string *)
let rec types2str = function
  | [] -> ""
  | [ctype] -> type2str ctype
  | ctype :: tail -> type2str ctype ^ " or " ^ (types2str tail)

let array_depth = function
  | ArrayDims (_, dims) -> List.length dims
  | _                   -> raise InvalidNode

let spec = function
  | ArrayDims (ctype, dims) -> (ctype, List.length dims)
  | ctype                   -> (ctype, 0)

let type2str_error = function
  | ArrayDims (ctype, dims) ->
    type2str ctype ^ "[" ^ repeat "," (List.length dims - 1) ^ "]"
  | ctype ->
    type2str ctype

let check_type ?(msg="") expected node =
  let got = typeof node in
  if expected <> Unknown && got <> Unknown && (spec got) <> (spec expected)
    then begin
      let msg = match msg with
      | "" -> sprintf "type mismatch: expected type %s, got %s"
              (type2str_error expected) (type2str_error got)
      | _ -> msg
      in
      [NodeMsg (node, msg)]
    end
    else []

let op_types = function
  | Not | And | Or                      -> [Bool]
  | Mod                                 -> [Int]
  | Neg | Sub | Div | Lt | Le | Gt | Ge -> [Int; Float]
  | Add | Mul | Eq | Ne                 -> [Bool; Int; Float]

let op_result_type opnd_type = function
  | Not | And | Or | Eq | Ne | Lt | Le | Gt | Ge -> Bool
  | Neg | Add | Sub | Mul | Div | Mod            -> opnd_type

(* Check if the given operator can be applied to the given type *)
let check_type_op allowed_types desc node =
  let got = typeof node in
  if got <> Unknown && not (List.mem got allowed_types)
    then
      [NodeMsg (node, sprintf
        "%s cannot be applied to type %s, only to %s"
        desc (type2str got) (types2str allowed_types))]
    else
      []

let check_dims_match dims dec_type errnode =
  match (List.length dims, array_depth dec_type) with
  | (got, expected) when got != expected ->
    let msg = sprintf
      "dimension mismatch: expected %d indices, got %d"
      expected got
    in
    [NodeMsg (errnode, msg)]
  | _ -> []

let err_map f nodes =
  let n, e = List.split (List.map f nodes) in
  (n, List.concat e)

let default_unknown ctype = function [] -> ctype | _ -> Unknown

let rec typecheck node =
  let add_error node msg =
    let node, err = traverse_list typecheck node in
    (annotate (Type Unknown) node, NodeMsg (node, msg) :: err)
  in
  let check_trav ctype node =
    let node, err = typecheck node in
    (node, err @ check_type ctype node)
  in
  match node with
  | FunUse ((FunDec (ret_type, name, params, _) as dec), args, ann)
  | FunUse ((FunDef (_, ret_type, name, params, _, _) as dec), args, ann) ->
    begin
      match (List.length args, List.length params) with
      | (nargs, nparams) when nargs != nparams ->
        add_error node (sprintf
          "function \"%s\" expects %d arguments, got %d"
          name nparams nargs)
      | _ ->
        let args, aerr = err_map typecheck args in
        let check_arg_type arg param = check_type (typeof param) arg in
        let err = List.concat (List.map2 check_arg_type args params) in
        (FunUse (dec, args, Type ret_type :: ann), aerr @ err)
    end

  (* Operators match operand types and get a new type based on the operator *)
  | Monop (op, opnd, ann) ->
    let opnd, oerr = typecheck opnd in
    let desc = sprintf "unary operator \"%s\"" (op2str op) in
    let err = check_type_op (op_types op) desc opnd in
    let res_type = default_unknown (typeof opnd) err in
    (Monop (op, opnd, Type (op_result_type res_type op) :: ann),
     oerr @ err)

  | Binop (op, left, right, ann) ->
    let left, lerr = typecheck left in
    let right, rerr = typecheck right in
    let desc = sprintf "binary operator \"%s\"" (op2str op) in
    let err =
      (* Only compare operand types if left operand has a valid type *)
      match check_type_op (op_types op) desc left with
      | []  -> check_type (typeof left) right
      | err -> err
    in
    let res_type = default_unknown (typeof left) err in

    (* Check for division by zero *)
    begin
      match (op, right) with
      | (Div, Const (IntVal 0l, _)) ->
        node_warning right "division by zero"
      | _ -> ()
    end;

    (Binop (op, left, right, Type (op_result_type res_type op) :: ann),
     lerr @ rerr @ err)

  (* Conditions must be bool, and right-hand type must match left-hand type *)
  | Cond (cond, texpr, fexpr, ann) ->
    let cond, cerr = check_trav Bool cond in
    let texpr, terr = typecheck texpr in
    let fexpr, ferr = check_trav (typeof texpr) fexpr in
    (Cond (cond, texpr, fexpr, Type (typeof texpr) :: ann), cerr @ terr @ ferr)

  (* Only basic types can be typecasted *)
  | TypeCast (ctype, value, ann) ->
    let value, err = typecheck value in
    let err = err @ check_type_op [Bool; Int; Float] "typecast" value in
    (TypeCast (ctype, value, Type ctype :: ann), err)

  (* Array allocation dimensions must have type int *)
  | Allocate (dec, dims, ann) ->
    let dims, err = err_map typecheck dims in
    let err = err @ List.concat (List.map (check_type Int) dims) in
    (Allocate (dec, dims, ann), err)

  (* Array dimensions are always integers *)
  | Dim (name, ann) ->
    (Dim (name, Type Int :: ann), [])

  (* Void functions may have no return statement, other functions must have a
   * return statement of valid type *)
  | FunDef (export, ret_type, name, params, body, ann) ->
    let params, perr = err_map typecheck params in
    let body, berr = typecheck body in
    let rec find_return = function
      | []                         -> None
      | [Return (value, _) as ret] -> Some (ret, typeof value)
      | hd :: tl                   -> find_return tl
    in
    let err =
      match (ret_type, find_return (block_body body)) with
      | (Void, Some (ret, _)) ->
        [NodeMsg (ret, "void function should not have a return value")]

      | ((Bool | Int | Float), None) ->
        [NodeMsg (node, sprintf
          "expected return value of type %s for function \"%s\""
          (type2str ret_type) name)]

      | ((Bool | Int | Float), Some (ret, t)) when t != ret_type ->
        [NodeMsg (ret, sprintf
          "function \"%s\" has return type %s, got %s"
          name (type2str ret_type) (type2str t))]

      | _ -> []
    in
    (FunDef (export, ret_type, name, params, body, ann), perr @ berr @ err)

  (* Conditions in must have type bool *)
  | If (cond, body, ann) ->
    let cond, cerr = check_trav Bool cond in
    let body, berr = typecheck body in
    (If (cond, body, ann), cerr @ berr)
  | IfElse (cond, tbody, fbody, ann) ->
    let cond, cerr = check_trav Bool cond in
    let tbody, terr = typecheck tbody in
    let fbody, ferr = typecheck fbody in
    (IfElse (cond, tbody, fbody, ann), cerr @ terr @ ferr)
  | While (cond, body, ann) ->
    let cond, cerr = check_trav Bool cond in
    let body, berr = typecheck body in
    (While (cond, body, ann), cerr @ berr)
  | DoWhile (cond, body, ann) ->
    let body, berr = typecheck body in
    let cond, cerr = check_trav Bool cond in
    (DoWhile (cond, body, ann), berr @ cerr)

  (* Constants *)
  | Const (BoolVal value, ann) ->
    (Const (BoolVal value, Type Bool :: ann), [])
  | Const (IntVal value, ann) ->
    (Const (IntVal value, Type Int :: ann), [])
  | Const (FloatVal value, ann) ->
    (Const (FloatVal value, Type Float :: ann), [])

  (* Variables inherit the type of their declaration *)
  | VarUse (dec, None, ann) ->
    (VarUse (dec, None, Type (typeof dec) :: ann), [])

  | VarUse (dec, Some dims, ann) ->
    (* Dimensions must have int type *)
    let dims, err = err_map typecheck dims in
    let err = err @ List.concat (List.map (check_type Int) dims) in

    (* Number of indices must match number of array dimensions *)
    let err = err @ check_dims_match dims (typeof dec) node in
    (VarUse (dec, Some dims, Type (basetypeof dec) :: ann), err)

  (* Array pointers cannot be re-assigned, because array dimension reduction
   * makes assumptions about dimensions of an array *)
  | VarLet (dec, None, _, _) when is_array dec ->
    add_error node "cannot assign value to array pointer after initialisation"

  (* Assigned values must match variable declaration *)
  | VarLet (dec, None, value, ann) ->
    let value, err = typecheck value in
    let err = err @ check_type (typeof dec) value in
    (VarLet (dec, None, value, ann), err)

  | VarLet (dec, Some dims, value, ann) ->
    (* Number of indices must match number of array dimensions *)
    let err1 = check_dims_match dims (typeof dec) node in

    (* Array indices must be ints *)
    let dims, err2 = err_map typecheck dims in
    let err2 = err2 @ List.concat (List.map (check_type Int) dims) in

    (* Assigned value must match array base type *)
    let value, err3 = typecheck value in
    let err3 = err3 @ check_type (basetypeof dec) value in

    (VarLet (dec, Some dims, value, ann), err1 @ err2 @ err3)

  (* ArrayConst initialisations are transformed during desugaring, so any
   * occurrences that are left are illegal *)
  | ArrayConst _ ->
    add_error node "array constants can only be used in array initialisation"

  | _ -> traverse_list typecheck node

let phase = function
  | Ast node ->
    let node, err = typecheck node in
    Ast (quit_on_error node err)
  | _ -> raise InvalidInput