civicaml/phases/index.ml

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")