open Types
open Util

let ret_incr r = let cur = !r in incr r; cur

(* Tag declarations with stack frame indices *)
let tag_index program =
  let nglobs = ref 0 in
  let nimportvar = ref 0 in
  let nimportfun = ref 0 in
  let consts = Hashtbl.create 32 in
  let rec trav_localfuns trav = function
    | LocalFuns body -> LocalFuns (List.map trav body)
    | node -> traverse_unit (trav_localfuns trav) node
  in
  let rec tag stacklen callstack node =
    let trav = tag stacklen callstack in
    match node with
    | GlobalDef _ ->
      let index = ret_incr nglobs in
      annotate (Index index) (traverse_unit trav node)

    | FunDef (export, rtype, name, params, body, ann) ->
      (* label name for local function is "__<parent_label>_<name>" *)
      let callstack = name :: callstack in
      let prefix = if List.length callstack > 1 then "__" else "" in
      let label = prefix ^ String.concat "_" (List.rev callstack) in

      let stacklen = ref 0 in
      let trav = tag stacklen callstack in

      (* Traverse own function body first *)
      let params = List.map trav params in
      let body = trav body in
      let ann = Index (!stacklen - List.length params) :: ann in

      (* Traverse local functions after the function body *)
      let body = trav_localfuns trav body in

      FunDef (export, rtype, name, params, body, LabelName label :: ann)

    | LocalFuns _ -> node

    | VarDec _ | Param _ | Dim _ ->
      let index = ret_incr stacklen in
      annotate (Index index) (traverse_unit trav node)

    | GlobalDec (_, name, _) ->
      let index = ret_incr nimportvar in
      annotate (Index index) node

    | FunDec (_, name, _, _) ->
      let index = ret_incr nimportfun in
      annotate (LabelName name) (annotate (Index index) node)

    | Const (value, _) when not (is_immediate_const value) ->
      let index =
        if Hashtbl.mem consts value then
          Hashtbl.find consts value
        else
          let index = Hashtbl.length consts in
          Hashtbl.add consts value index;
          index
      in
      annotate (Index index) node

    | _ -> traverse_unit trav node
  in tag (ref 0) [] program

(* Undo context analysis *)
let rec strip_context = function
  | VarUse (dec, dims, ann) ->
    Var (nameof dec, optmap strip_context dims, ann)

  | VarLet (dec, dims, value, ann) ->
    Assign (nameof dec, optmap strip_context dims, strip_context value, ann)

  | FunUse (dec, args, ann) ->
    FunCall (nameof dec, List.map strip_context args, ann)

  | node -> traverse_unit strip_context node

let phase = function
  | Ast node ->
    let tagged = tag_index (strip_context node) in
    Ast (Context.analyse_context tagged)
  | _ -> raise InvalidInput