open Printf
open Types
open Util

type nametype = Varname of string | Funcname of string

let type2str = function Funcname _ -> "function" | Varname _ -> "variable"

let mapfind name tbl =
  if Hashtbl.mem tbl name then Some (Hashtbl.find tbl name) else None

let check_in_scope name errnode scope err =
  let vars, funs = scope in
  let name, tbl, other_map, desired_type = match name with
  | Varname  name -> (name, vars, funs, "variable")
  | Funcname name -> (name, funs, vars, "function")
  in
  match mapfind name tbl with
  | Some (dec, dec_depth, _) ->
    (dec, dec_depth)
  | None ->
    let msg = match mapfind name other_map with
    | Some _ -> sprintf "\"%s\" is not a %s" name desired_type
    | None   -> sprintf "undefined %s \"%s\"" desired_type name
    in
    err := !err @ [NodeMsg (errnode, msg)];
    (DummyNode, -1)

let add_to_scope name dec depth (vars, funs) =
  let name, tbl, name_type = match name with
  | Varname name  -> name, vars, "variable"
  | Funcname name -> name, funs, "function"
  in
  match mapfind name tbl with
  (* Identifiers of lower depth may be overwritten, but idenetifiers at
    * the same depth must be unique for consistency *)
  | Some (orig, orig_depth, _) when orig_depth >= depth ->
    (* For generated variables, don't gove an error, since the error variable
     * is a derived array dimension of a redefined array, which will yield an
     * error later on *)
    if is_generated_id name then
      Hashtbl.replace tbl name (dec, depth, name_type)
    else
      let msg = sprintf "Error: cannot redeclare %s \"%s\"" name_type name in
      prerr_loc_msg (locof dec) msg;
      prerr_loc_msg (locof orig) "Previously declared here:";
      raise (FatalError NoMsg)
  | Some _ ->
    Hashtbl.replace tbl name (dec, depth, name_type)
  | None ->
    Hashtbl.add tbl name (dec, depth, name_type)

let rec analyse scope depth node err =
  let rec collect node = match node with
    (* For extern array declarations, add the dimension names as well *)
    | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
      let t = ArrayDims (ctype, List.map (annotate (Depth depth)) dims) in
      let node = GlobalDec (t, name, Depth depth :: ann) in
      add_to_scope (Varname name) node depth scope;
      node

    (* For variables, add the name (array dimensions are added
     * implicitly, since they have separate VarDec nodes which were added
     * during the desugaring phase *)
    | VarDec (_, name, _, _)
    | GlobalDec (_, name, _)
    | GlobalDef (_, _, name, _, _) ->
      let node = annotate (Depth depth) node in
      add_to_scope (Varname name) node depth scope;
      node

    (* Functions are traversed later on, for now only add the name *)
    | FunDec (_, name, _, _)
    | FunDef (_, _, name, _, _, _) ->
      let node = annotate (Depth depth) node in
      add_to_scope (Funcname name) node depth scope;
      node

    (* For a variable or function call, look for its declaration in the
     * current scope and save a its type/depth information  *)
    | Var (name, dims, ann) ->
      let dec, dec_depth = check_in_scope (Varname name) node scope err in
      VarUse (dec, optmap collect dims, Depth depth :: ann)

    | FunCall (name, args, ann) ->
      let dec, dec_depth = check_in_scope (Funcname name) node scope err in
      FunUse (dec, List.map collect args, Depth depth :: ann)

    (* Assign statements are replaced with VarLet nodes, which stores the
     * declaration of the assigned variable *)
    | Assign (name, dims, value, ann) ->
      let dec, dec_depth = check_in_scope (Varname name) node scope err in
      VarLet (dec, optmap collect dims, collect value, Depth depth :: ann)

    | Allocate (dec, dims, ann) ->
      let dec, dec_depth = check_in_scope (Varname (nameof dec)) node scope err in
      Allocate (dec, List.map collect dims, Depth depth :: ann)

    | _ -> traverse_unit collect node
  in

  let rec traverse scope depth = function
    (* Increase nesting level when entering function *)
    | FunDef (export, ret_type, name, params, body, ann) ->
      let vars, funs = scope in
      let local_scope = (Hashtbl.copy vars, Hashtbl.copy funs) in
      let params = List.map (traverse local_scope (depth + 1)) params in
      let body = analyse local_scope (depth + 1) body err in
      FunDef (export, ret_type, name, params, body, ann)

    | Param (ArrayDims (ctype, dims), name, ann) ->
      let rec add_dims = function
        | [] -> []
        | Dim (name, ann) :: tl ->
          let dim = Dim (name, Depth depth :: ann) in
          add_to_scope (Varname name) dim depth scope;
          dim :: (add_dims tl)
        | _ -> raise InvalidNode
      in
      let node = Param (ArrayDims (ctype, add_dims dims), name, ann) in
      add_to_scope (Varname name) node depth scope;
      node

    | VarDec _ as node -> node

    | Param (_, name, _) as node ->
      let node = annotate (Depth depth) node in
      add_to_scope (Varname name) node depth scope;
      node

    (* Do not traverse into external function declarations, since their
     * parameters must not be added to the namespace *)
    | FunDec _ as node -> node

    | node  -> traverse_unit (traverse scope depth) node
  in

  (*
   * First collect all definitions at the current depth. Then, traverse into
   * functions with a copy of the current scope. This is needed because
   * functions can access all identifiers in their surrounding scope.
   * E.g., the following is allowed:
   *
   * void foo() { glob = 1; }
   * int glob;
   *)
  collect node |> traverse scope depth

let analyse_context program =
  let scope = (Hashtbl.create 20, Hashtbl.create 20) in
  let err = ref [] in
  let node = analyse scope 0 program err in
  quit_on_error node !err

let phase = function
  | Ast node -> Ast (analyse_context node)
  | _ -> raise InvalidInput