open Types
open Util

let assemble program =
    let labcounter = ref 0 in
    let genlabel suffix =
        labcounter := !labcounter + 1;
        string_of_int !labcounter ^ "_" ^ suffix
    in

    let consts = Hashtbl.create 20 in

    let rec trav_args callstack localfuns node =
        let trav = trav_args callstack localfuns in
        let rec trav_all = function
            | [] -> []
            | hd :: tl -> trav hd @ (trav_all tl)
        in
        let rec traverse_localfuns = function
            | LocalFuns funs -> trav_all funs
            | Block body -> List.concat (List.map traverse_localfuns body)
            | _ -> []
        in
        match node with
        | Program (decls, _) ->
            trav_all decls

        | FunDec (ret_type, name, params, _) ->
            [Import (name, ret_type, List.map typeof params)]

        | FunDef (export, ret_type, name, params, body, _) ->
            localfuns := node :: !localfuns;
            let callstack = name :: callstack in
            let localfuns = ref [] in
            let label = String.concat "$" (List.rev callstack) in
            (if export then
                let param_types = List.map typeof params in
                [Export (name, ret_type, param_types, label)]
            else []) @
            [
                Comment ("function \"" ^ label ^ "\":");
                Label label;
                RtnEnter (indexof node);
            ] @
            (trav_args callstack localfuns body) @
            (match ret_type with Void -> [Ret Void] | _ -> []) @
            [EmptyLine] @
            (traverse_localfuns body)

        (* Local fucntions are traversed elsewhere *)
        | LocalFuns _ -> []

        | Block body -> trav_all body

        | VarLet (dec, None, value, _) ->
            let store = match (depthof dec, depthof node) with
                | (0, _)            -> StoreGlob (typeof dec, indexof dec)
                | (a, b) when a = b -> StoreLoc  (typeof dec, indexof dec)
                | (a, b)            -> StoreRel  (typeof dec, b - a, indexof dec)
            in
            trav value @ [store]

        | Return (value, _) ->
            trav value @ [Ret (typeof value)]

        | Const (BoolVal _, _) ->
            [LoadImm node]
        | Const (value, _) ->
            let def = if Hashtbl.mem consts value then [] else (
                Hashtbl.add consts value true;
                [ConstDef value]
            ) in
            def @ [LoadConst (typeof node, indexof node)]

        | _ -> []
        (*| _ -> raise InvalidNode*)
    in
    trav_args [] (ref []) program

let rec phase input =
    log_line 2 "- Assembly";
    match input with
    | Ast node -> Assembly (assemble node)
    | _ -> raise (InvalidInput "assembly")