open Printf
open Types
open Util

let rec add_depth depth node =
  match node with
  | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
    (* XXX: traversal should traverse into Dim nodes *)
    let dims = List.map (add_depth depth) dims in
    GlobalDec (ArrayDims (ctype, dims), name, Depth depth :: ann)
  | Param (ArrayDims (ctype, dims), name, ann) ->
    let dims = List.map (add_depth depth) dims in
    Param (ArrayDims (ctype, dims), name, Depth depth :: ann)
  | GlobalDec _
  | GlobalDef _
  | Param _
  | Dim _
  | VarDec _
  | Var _
  | FunCall _
  | Assign _
  | For _
  | VarUse _
  | FunUse _
  | VarLet _ ->
    annotate (Depth depth) node |> traverse_unit (add_depth depth)
  | FunDec _ ->
    annotate (Depth depth) node
  | FunDef (export, ret_type, name, params, body, ann) ->
    let params = List.map (add_depth (depth + 1)) params in
    let body = add_depth (depth + 1) body in
    FunDef (export, ret_type, name, params, body, Depth depth :: ann)
  | _ ->
    traverse_unit (add_depth depth) node

type identifier_type = Funcname | Varname

let typename = function Varname -> "variable" | Funcname -> "function"

let tblfind tbl name = try Some (Hashtbl.find tbl name) with Not_found -> None

let add_to_scope scope name dec namety err =
  match tblfind scope name with
  | Some orig when depthof orig >= depthof dec ->
    err := NodeMsg (orig, "\rPreviously declared here:") ::
           NodeMsg (dec, sprintf "Error: cannot redeclare %s \"%s\""
                         (typename namety) name) :: !err
  | Some _ ->
    Hashtbl.replace scope name dec
  | None ->
    Hashtbl.add scope name dec

let check_in_scope scope name errnode namety err =
  match tblfind scope name with
  | Some dec -> dec
  | None ->
    let msg = sprintf "undefined %s \"%s\"" (typename namety) name in
    err := NodeMsg (errnode, msg) :: !err;
    DummyNode

let prt_vars vars =
  let prt name dec = prerr_string (name ^ ", ") in
  Hashtbl.iter prt vars;
  prerr_endline "(end)"

let analyse do_rename program =
  let err = ref [] in

  (* Add functions at the current depth to the function scope, do not traverse
   * into nested functions *)
  let rec collect_funs funs node =
    match node with
    | FunDec (_, name, _, _)
    | FunDef (_, _, name, _, _, _) ->
      (* TODO: don't copy function body to save memory *)
      add_to_scope funs name node Funcname err;
      node
    | _ -> traverse_unit (collect_funs funs) node
  in

  (* Traverse through statements in the current scope, checking and replacing
   * variable occurrences. Add newly declared variables to the variable scope
   * on-the-fly. *)
  let rec traverse scope node =
    let trav = traverse scope in
    let trav_dims = function
      | ArrayDims (ctype, dims) ->
        ArrayDims (ctype, List.map trav dims)
      | ctype -> ctype
    in
    let vars, funs, repl = scope in
    let check_rename node rename =
      let name = nameof node in
      let shadows_higher_scope =
        do_rename && (Hashtbl.mem vars name || Hashtbl.mem repl name)
      in

      (* Trigger duplication error and make sure following duplication errors
       * will refer to a non-generated variable name *)
      add_to_scope vars name node Varname err;

      if shadows_higher_scope then begin
        let newname = fresh_id name in
        Hashtbl.replace repl name newname;
        let newnode = rename newname in
        add_to_scope vars newname newnode Varname err;
        newnode
      end else
        node
    in
    let add_dims = function
      | ArrayDims (ctype, dims) ->
        let rec add = function
          | [] -> []
          | (Dim (name, ann) as dim) :: tl ->
            check_rename dim (fun name -> Dim (name, ann)) :: add tl
          | _ -> raise InvalidNode
        in
        ArrayDims (ctype, add dims)
      | ctype -> ctype
    in
    match node with
    | Program (decls, ann) ->
      Program (Block decls |> trav |> trav_funs scope |> block_body, ann)

    | FunDec _ | FunDef _ -> node

    | GlobalDef (export, ctype, name, init, ann) ->
      let node = GlobalDef (export, trav_dims ctype, name, optdo trav init, ann) in
      add_to_scope vars name node Varname err;
      node

    | VarDec (ctype, name, init, ann) ->
      let ctype = trav_dims ctype in
      let init = optdo trav init in
      let node = VarDec (ctype, name, init, ann) in
      check_rename node (fun name -> VarDec (ctype, name, init, ann))

    | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
      let rec replace_dims i = function
        | [] -> []
        | Dim (dimname, ann) :: tl ->
          let newname = generate_array_dim name i in
          Hashtbl.add repl dimname newname;
          Dim (newname, ann) :: replace_dims (i + 1) tl
        | _ -> raise InvalidNode
      in
      let ctype = ArrayDims (ctype, replace_dims 0 dims) in
      let node = GlobalDec (add_dims ctype, name, ann) in
      add_to_scope vars name node Varname err;
      node

    | GlobalDec (ctype, name, ann) ->
      add_to_scope vars name node Varname err;
      node

    | Param (ctype, name, ann) ->
      let ctype = add_dims ctype in
      let node = Param (ctype, name, ann) in
      check_rename node (fun name -> Param (ctype, name, ann))

    | For (name, start, stop, step, body, ann) ->
      let start, stop, step = trav start, trav stop, trav step in

      (* For-loops are a special case: the loop counter defines a new scope
       * which is allowed to shadow existing local variables, and allows for
       * nested loops with the same induction variables. Replace the loop
       * counter with a fresh variable to enforce this behaviour, and avoid
       * having to replace variables during desigaring. *)
      (* FIXME: only create new variable if necessary *)
      let newname = fresh_id name in
      let node = For (newname, start, stop, step, body, ann) in
      Hashtbl.add vars name node;
      Hashtbl.add vars newname node;
      Hashtbl.add repl name newname;

      let node = For (newname, start, stop, step, trav body, ann) in
      Hashtbl.remove repl name;
      Hashtbl.remove vars newname;
      Hashtbl.remove vars name;
      node

    (* Perform renaming *)
    | Var (name, dims, ann) when Hashtbl.mem repl name ->
      trav (Var (Hashtbl.find repl name, dims, ann))

    | Assign (name, dims, value, ann) when Hashtbl.mem repl name ->
      trav (Assign (Hashtbl.find repl name, dims, value, ann) )

    (* Replace variables or function calls with use-nodes which contain the
     * entire declaration   *)
    | Var (name, dims, ann) ->
      let dec = check_in_scope vars name node Varname err in
      VarUse (dec, optmap trav dims, ann)

    | FunCall (name, args, ann) ->
      let dec = check_in_scope funs name node Funcname err in
      FunUse (dec, List.map trav args, ann)

    (* Assign statements are replaced with VarLet nodes, which stores the
     * declaration of the assigned variable *)
    | Assign (name, dims, value, ann) ->
      begin
        match check_in_scope vars name node Varname err with
        | For _ ->
          err := NodeMsg (node, "cannot assign to induction variable") :: !err;
          node
        | dec ->
          VarLet (dec, optmap trav dims, trav value, ann)
      end

    (* Also support intermediary nodes because context analysis is re-run later
     * on to propagate new declaration properties *)
    | VarUse (dec, dims, ann) ->
      VarUse (Hashtbl.find vars (nameof dec), optmap trav dims, ann)

    | FunUse (dec, args, ann) ->
      FunUse (Hashtbl.find funs (nameof dec), List.map trav args, ann)

    | VarLet (dec, dims, value, ann) ->
      VarLet (Hashtbl.find vars (nameof dec), optmap trav dims, trav value, ann)

    | Allocate (dec, dims, ann) ->
      Allocate (Hashtbl.find vars (nameof dec), List.map trav dims, ann)

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

    | _ -> traverse_unit trav node

  and trav_funs scope = function
    (* Copy scope when entering a function body *)
    | FunDef (export, ret_type, name, params, body, ann) ->
      let vars, funs, repl = scope in
      let locfuns = Hashtbl.copy funs in
      let locscope = (Hashtbl.copy vars, locfuns, Hashtbl.copy repl) in
      let params = List.map (traverse locscope) params in
      let body =
        collect_funs locfuns body |> traverse locscope |> trav_funs locscope
      in
      FunDef (export, ret_type, name, params, body, ann)

    | node -> traverse_unit (trav_funs scope) node
  in

  let vars = Hashtbl.create 32 in
  let funs = Hashtbl.create 16 in
  let repl = Hashtbl.create 16 in
  add_depth 0 program |>
  collect_funs funs |>
  traverse (vars, funs, repl) |>
  quit_on_error (List.rev !err)

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