open Types
open Util

let create_param ctype name =
  let param = Param (ctype, name, [Depth 1]) in
  let value = VarUse (param, None, [Type ctype; Depth 1]) in
  (param, value)

let call node args depth =
  match node with
  | FunDec (ctype, name, _, _) as dec ->
    FunUse (dec, args, [Type ctype; Depth depth])
  | _ -> raise InvalidNode

let create_getset globals = function
  | GlobalDef (true, ArrayDims (ctype, _), name, None, ann) as dec ->
    (* Getters for array variable: create getter for given index Note that
     * getters and setters for dimensions are automatically generated,
     * because they have been put into new global variables during the
     * desugaring phase *)
    let (param, index) = create_param Int (fresh_var "index") in
    let var = VarUse (dec, Some [index], [Type ctype; Depth 1]) in
    let body = Block [Return (var, [])] in
    let getter = FunDef (true, ctype, name ^ "$get", [param], body, []) in

    (* Setters for array variable: create setter for given index *)
    let (param1, index) = create_param Int (fresh_var "index") in
    let (param2, value) = create_param ctype (fresh_var "value") in
    let body = Block [VarLet (dec, Some [index], value, [])] in
    let setter = FunDef (true, Void, name ^ "$set", [param1; param2], body, []) in

    [getter; setter]

  | GlobalDef (true, ctype, name, None, ann) as dec ->
    (* Getter for basic variable type: return the variable *)
    let var = VarUse (dec, None, [Type ctype; Depth 1]) in
    let body = [Return (var, [])] in
    let getter = FunDef (true, ctype, name ^ "$get", [], Block body, []) in

    (* Setter for basic variable type: assign the variable *)
    let (param, value) = create_param ctype (fresh_var "value") in
    let body = [VarLet (dec, None, value, [])] in
    let setter = FunDef (true, Void, name ^ "$set", [param], Block body, []) in

    [getter; setter]

  | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
    (* External array variable: create getter and setter for a given index. Now
     * we also need to generate functions for dimensions since they are NOT
     * added as new variables during desugaring. *)
    let rec add_dims i = function
      | [] -> []
      | Dim (dimname, ann) :: tl ->
        let newname = name ^ "$" ^ string_of_int i in

        let getter = FunDec (ctype, newname ^ "$get", [], []) in

        let (param, _) = create_param ctype "value" in
        let setter = FunDec (Void, newname ^ "$set", [param], []) in

        Hashtbl.add globals dimname (call getter, call setter);
        getter :: setter :: (add_dims (i + 1) tl)
      | _ -> raise InvalidNode
    in
    let dimfuncs = add_dims 1 dims in

    let (param, _) = create_param Int "index" in
    let getter = FunDec (ctype, name ^ "$get", [param], []) in

    let (param1, index) = create_param Int "index" in
    let (param2, value) = create_param ctype "value" in
    let setter = FunDec (Void, name ^ "$set", [param1; param2], []) in

    Hashtbl.add globals name (call getter, call setter);
    getter :: setter :: dimfuncs

  (* Getter for basic variable type: return the variable *)
  | GlobalDec (ctype, name, ann) ->
    let getter = FunDec (ctype, name ^ "$get", [], []) in

    let (param, _) = create_param ctype "value" in
    let setter = FunDec (Void, name ^ "$set", [param], []) in

    Hashtbl.add globals name (call getter, call setter);
    [getter; setter]

  | _ -> raise InvalidNode

(* Create getter/setter functions for exported/imported variables *)
let rec create_funcs globals = function
  | Program (decls, ann) ->
    let decls = List.map (create_funcs globals) decls in
    Program (flatten_blocks (List.map (create_funcs globals) decls), ann)

  | GlobalDef (true, ctype, name, None, ann) as node ->
    Block (GlobalDef (false, ctype, name, None, ann) ::
    (create_getset globals node))

  | GlobalDec (ctype, name, ann) as node ->
    Block (create_getset globals node)

  | node -> transform_children (create_funcs globals) node

(* Replace uses for imported/exported variabels with getter/setter functions *)
let rec replace_vars scope depth = function
  (* Variable names may be redefined in function scopes *)
  | (VarDec (_, name, _, _) as node)
  | (Param (_, name, _) as node) when Hashtbl.mem scope name ->
    Hashtbl.remove scope name;
    node

  (* Copy scope when traversing into function,, and restore afterwards *)
  | FunDef (export, ret_type, name, params, body, ann) ->
    let local_scope = Hashtbl.copy scope in
    let trav = replace_vars local_scope (depth + 1) in
    let params = List.map trav params in
    FunDef (export, ret_type, name, params, trav body, ann)

  (* Use of regular external variable *)
  | VarUse (dec, None, _) when Hashtbl.mem scope (nameof dec) ->
    let (get, _) = Hashtbl.find scope (nameof dec) in
    get [] depth

  (* Dereference of external array *)
  | VarUse (dec, Some indices, _) when Hashtbl.mem scope (nameof dec) ->
    let (get, _) = Hashtbl.find scope (nameof dec) in
    get indices depth

  | VarLet (dec, dims, value, _) when Hashtbl.mem scope (nameof dec) ->
    let dims = optmapl (replace_vars scope depth) dims in
    let (_, set) = Hashtbl.find scope (nameof dec) in
    Expr (set (dims @ [replace_vars scope depth value]) depth)

  | node -> transform_children (replace_vars scope depth) node

let phase = function
  | Ast node ->
    let globals = Hashtbl.create 20 in
    let node = create_funcs globals node in
    Ast (replace_vars globals 0 node)
  | _ -> raise (InvalidInput "extern vars")