open Types
open Util

let tag_index program =
    let nglobs = ref 0 in
    let consts = Hashtbl.create 32 in
    let rec trav stacklen node = match node with
        | GlobalDef _ ->
            let index = !nglobs in
            nglobs := !nglobs + 1;
            annotate (Index index) (transform_children (trav stacklen) node)

        | FunDef _ ->
            let stacklen = ref 0 in
            let node = transform_children (trav stacklen) node in
            annotate (Index !stacklen) node

        | VarDec _ | Dim _ ->
            let index = !stacklen in
            stacklen := !stacklen + 1;
            annotate (Index index) (transform_children (trav stacklen) node)

        | 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 stacklen) node
    in trav (ref 0) program

let rec strip = function
    | VarUse (dec, dims, ann) ->
        Var (nameof dec, optmap strip dims, ann)

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

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

    | node -> transform_children strip node

let rec phase input =
    log_line 1 "- Depth analysis";
    match input with
    | Ast node -> Ast (Context_analysis.analyse_context (tag_index (strip node)))
    | _ -> raise (InvalidInput "depth analysis")