civicaml/phases/assemble.ml

open Printf
open Types
open Util
open Stringify

let comline comment = InlineComment (EmptyLine, comment)

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 node =
        let rec trav_all = function
            | [] -> []
            | hd :: tl -> trav hd @ (trav_all tl)
        in
        let rec traverse_localfuns = function
            | LocalFuns body -> List.concat (List.map trav body)
            | Block body -> List.concat (List.map traverse_localfuns body)
            | _ -> []
        in
        match node with
        (* Global *)
        | Program (decls, _) ->
            trav_all decls

        | GlobalDef (_, ctype, name, _, _) ->
            [Comment (sprintf "global var \"%s\" at index %d" name (indexof node));
             Global ctype]

        | FunDec (ret_type, name, params, _) ->
            [Comment (sprintf "extern fun \"%s\" at index %d" name (indexof node));
             Import (name, ret_type, List.map typeof params)]

        | FunDef (export, ret_type, name, params, body, _) ->
            let label = labelof node in
            (if export then
                let param_types = List.map typeof params in
                [Export (name, ret_type, param_types, label)]
            else []) @
            [
                Comment (sprintf "fun \"%s\" with %d local vars" label (indexof node));
                Label label;
                RtnEnter (indexof node);
            ] @
            (trav body) @
            (match ret_type with Void -> [Ret Void] | _ -> []) @
            [EmptyLine] @
            (traverse_localfuns body)

        | VarDec (_, name, _, _) ->
            [comline (sprintf "local var \"%s\" at index %d" name (indexof node))]

        | LocalFuns _ -> []

        | Block body | VarDecs body -> trav_all body

        (* Statements *)
        | VarLet (dec, None, value, _) ->
            let store = match (depthof dec, depthof node) with
                | (0, _)            -> Store (typeof dec, Glob,        indexof dec)
                | (a, b) when a = b -> Store (typeof dec, Current,     indexof dec)
                | (a, b)            -> Store (typeof dec, Rel (b - a), indexof dec)
            in
            trav value @ [InlineComment (store, node2str node)]

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

        | If (cond, body, _) ->
            let endlabel = genlabel "end" in
            (trav cond) @
            [Branch (false, endlabel);
             comline ("if (" ^ (node2str cond) ^ ") {")] @
            (trav body) @
            [comline "}";
             Label endlabel]

        | IfElse (cond, true_body, false_body, _) ->
            let elselabel = genlabel "else" in
            let endlabel = genlabel "end" in
            (trav cond) @
            [Branch (false, elselabel);
             comline ("if (" ^ (node2str cond) ^ ") {")] @
            (trav true_body) @
            [Jump endlabel;
             comline "} else {";
             Label elselabel] @
            (trav false_body) @
            [comline "}";
             Label endlabel]

        | While (cond, body, _) ->
            let startlabel = genlabel "while" in
            let endlabel = genlabel "end" in
            let com = ("while (" ^ (node2str cond) ^ ") {") in
            [Label startlabel] @
            (trav cond) @
            [InlineComment (Branch (false, endlabel), com)] @
            (trav body) @
            [Jump startlabel;
             Label endlabel;
             comline "}"]

        | DoWhile (cond, body, _) ->
            let startlabel = genlabel "dowhile" in
            let com = ("} while (" ^ (node2str cond) ^ ");") in
            [comline "do {";
             Label startlabel] @
            (trav body) @
            (trav cond) @
            [InlineComment (Branch (true, startlabel), com)]

        (* Expression statement pops the disregarded expression value from the
         * stack, if any *)
        | Expr value ->
            let pop = match typeof value with
                | Void -> [comline (node2str node)]
                | ctype -> [InlineComment (Pop ctype, node2str node)]
            in
            (trav value) @ pop

        (* Expressions *)
        | Const (BoolVal _ as value, _) ->
            [LoadImm value]

        | Const (value, _) ->
            Hashtbl.replace consts value (typeof node, indexof node);
            let load = LoadConst (typeof node, indexof node) in
            [InlineComment (load, node2str node)]

        | VarUse (dec, None, _) ->
            let load = match (depthof dec, depthof node) with
                | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
                | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
                | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
            in
            [InlineComment (load, node2str node)]

        | Monop (op, value, _) ->
            (trav value) @
            [InlineComment (Op (op, typeof value), op2str op)]

        | Binop (op, left, right, _) ->
            (trav left) @
            (trav right) @
            [InlineComment (Op (op, typeof left), op2str op)]

        | TypeCast (ctype, value, _) ->
            let vtype = typeof value in
            (match (ctype, vtype) with
            | (Float, Int) | (Int, Float) -> ()
            | _ -> raise (NodeError (node, "invalid typecast"))
            );
            trav value @ [Convert (vtype, ctype)]

        (* Function calls *)
        | FunUse (dec, args, _) ->
            let init = match (depthof dec, depthof node) with
                | (0, _)                -> RtnInit Glob
                | (a, b) when a = b     -> RtnInit Current
                | (a, b) when a = b + 1 -> RtnInit Local
                | (a, b)                -> RtnInit (Rel (b - a))
            in
            let jmp = match dec with
                | FunDec _ -> RtnJmp (ExternFun (indexof dec))
                | FunDef _ -> RtnJmp (LocalFun (List.length args, labelof dec))
                | _ -> raise InvalidNode
            in
            init :: (trav_all args) @ [jmp]

        | Arg value -> trav value

        (* Conditional expression (short-circuit evaluation) *)
        (*     <cond>
         *     branch_f else
         *     <true_expr>
         *     jump end
         * else:
         *     <false_expr>
         * end:
         *)
        | Cond (cond, texp, fexp, _) ->
            let elselabel = genlabel "false_expr" in
            let endlabel = genlabel "end" in
            (trav cond) @
            [Branch (false, elselabel)] @
            (trav texp) @
            [Jump (endlabel);
             Label (elselabel)] @
            (trav fexp) @
            [InlineComment (Label (endlabel), node2str node)]

        (* Arrays *)
        | Allocate (dec, dims, _) ->
            let store = match (depthof dec, depthof node) with
                | (0, _)            -> Store (typeof dec, Glob,    indexof dec)
                | (a, b) when a = b -> Store (typeof dec, Current, indexof dec)
                | _                 -> raise InvalidNode
            in
            trav_all dims @
            [NewArray (basetypeof dec, List.length dims);
             InlineComment (store, node2str node)]

        | VarUse (dec, Some dims, _) ->
            let load = match (depthof dec, depthof node) with
                | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
                | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
                | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
            in
            (trav_all dims) @                     (* push dimensions *)
            [InlineComment (load, nameof dec)] @  (* push array reference *)
            [InlineComment (LoadArray (basetypeof dec), node2str node)]

        | VarLet (dec, Some dims, value, _) ->
            let load = match (depthof dec, depthof node) with
                | (0, _)            -> Load (typeof dec, Glob,        indexof dec)
                | (a, b) when a = b -> Load (typeof dec, Current,     indexof dec)
                | (a, b)            -> Load (typeof dec, Rel (b - a), indexof dec)
            in
            (trav value) @                        (* push value *)
            (trav_all dims) @                     (* push dimensions *)
            [InlineComment (load, nameof dec)] @  (* push array reference *)
            [InlineComment (StoreArray (basetypeof dec), node2str node)]

        | _ -> [Comment ("FIXME: " ^ Stringify.node2str node)]
        (*| _ -> raise InvalidNode*)
    in
    let instrs = trav program in

    (* Sort aggregated constants and add definitions
     * If possible, this should be rewritten in the future because it's a little
     * cumbersome right now... *)
    let pairs = ref [] in
    let add_pair value (ctype, index) =
        let com = sprintf "index %d" index in
        pairs := (InlineComment (ConstDef (ctype, value), com), index) :: !pairs;
    in
    Hashtbl.iter add_pair consts;
    let cmp (_, i) (_, j) = compare i j in
    let sorted_pairs = List.sort cmp !pairs in
    let const_defs = List.map (fun (d, _) -> d) sorted_pairs in

    instrs @ const_defs

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