open Types
open Util
open Globals

let tag_index program =
  let nglobs = ref 0 in
  let nimport = ref 0 in
  let consts = Hashtbl.create 32 in
  let rec trav_localfuns trav = function
    | LocalFuns body -> LocalFuns (List.map trav body)
    | node -> transform_children (trav_localfuns trav) node
  in
  let rec tag stacklen callstack node =
    let trav = tag stacklen callstack in
    match node with
    | GlobalDef _ ->
      let index = !nglobs in
      nglobs := !nglobs + 1;
      annotate (Index index) (transform_children trav node)

    | FunDef (export, rtype, name, params, body, ann) ->
      (* label name for local function is "<parent_label>$<name>" *)
      let callstack = name :: callstack in
      let label = 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 = !stacklen in
      stacklen := !stacklen + 1;
      annotate (Index index) (transform_children trav node)

    | FunDec (_, name, _, _) ->
      let index = !nimport in
      nimport := !nimport + 1;
      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

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

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 -> transform_children strip_context node

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