civicaml/phases/assemble.ml

open Printf
open Types
open Util
open Stringify
open Globals

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 *)
  (* Immediate values are handled here, and not in the peephole optimizer,
   * for convenience: the indices in the constant table would be altered,
   * so entries cannot be removed. By this early detection (also during
   * index analysis), they are not added in the first place *)
  | Const (value, _) when is_immediate_const value ->
    [InlineComment (LoadImm value, node2str node)]

  | Const (value, _) ->
    Hashtbl.replace consts value (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) -> ()
    | _ ->
      let msg = sprintf
        "invalid typecast: %s -> %s"
        (type2str vtype) (type2str ctype)
      in raise (NodeError (node, msg))
    );
    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 - 1 -> RtnInit Current
    | (a, b) when a = b     -> RtnInit Local
    | (a, b)                -> RtnInit (Rel (b - a - 1))
    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 (List.rev 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 index =
  let com = sprintf "index %d" index in
  pairs := (InlineComment (ConstDef 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 phase = function
  | Ast node -> Assembly (assemble node)
  | _ -> raise (InvalidInput "assembly")