civicaml/phases/assemble.ml

open Printf
open Types
open Util
open Stringify

let get_scope node = function
  | GlobalDec _ -> Extern
  | dec ->
    match (depthof dec, depthof node) with
    | (0, _)            -> Glob
    | (a, b) when a = b -> Current
    | (a, b)            -> Rel (b - a)

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 (false, ctype, name, _, _) ->
      [Comment (sprintf "global var \"%s\" at index %d" name (indexof node));
       Global ctype]

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

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

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

    | FunDef (export, ret_type, name, params, body, _) ->
      let last_if_return body =
        match List.fold_left (fun _ x -> x) DummyNode (block_body body) with
        | Return _ as ret -> Some ret
        | _               -> None
      in
      let add_explicit_return body =
        match last_if_return body with
        | Some _ -> body
        | None   -> Block (block_body body @ [Return (None, [])])
      in
      let label = labelof node in
      begin
        if export
        then [ExportFun (name, ret_type, List.map typeof params, label)]
        else []
      end @
      [Comment (sprintf "function \"%s\" with %d parameters and %d local vars"
                        label (List.length params) (indexof node));
       Label label] @
      (match indexof node with 0 -> [] | n -> [RtnEnter n]) @
      (trav_all params) @
      (trav (add_explicit_return body)) @
      [EmptyLine] @
      (traverse_localfuns body)

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

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

    | LocalFuns _ -> []

    | Block body | VarDecs body -> trav_all body

    (* Statements *)

    | VarLet (dec, None, value, _) ->
      let store = Store (typeof dec, get_scope node dec, indexof dec) in
      trav value @ [InlineComment (store, node2str node)]

    | Return (None, _) ->
      [InlineComment (Ret Void, node2str node)]

    | Return (Some 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 -> []
      | ctype -> [InlineComment (Pop ctype, "disregard return value")]
      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 (GlobalDec (ctype, _, _) as dec, None, _) ->
      let load = Load (ctype, Extern, indexof dec) 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
      begin
        match (ctype, vtype) with
        | (Float, Int) | (Int, Float) -> ()
        | _ ->
          let msg = sprintf
            "invalid typecast: %s -> %s"
            (type2str vtype) (type2str ctype)
          in
          raise (FatalError (NodeMsg (node, msg)))
      end;
      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
      InlineComment (init, nameof dec) ::
      (trav_all args) @
      [InlineComment (jmp, node2str node)]

    | 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, [dim], _) ->
      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 dim @
      [NewArray (basetypeof dec);
       InlineComment (store, node2str node)]

    | Allocate _ ->
      raise (FatalError (NodeMsg (node, "invalid number of array dimensions \
                                         (should be one-dimensional)")))

    | VarUse (dec, Some dims, _) ->
      let load = Load (typeof dec, get_scope node dec, 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 = Load (typeof dec, get_scope node dec, indexof dec) in
      (trav_all (List.rev dims)) @          (* push dimensions *)
      (trav value) @                        (* push value *)
      (trav_all dims) @                     (* push dimensions *)
      [InlineComment (load, nameof dec)] @  (* push array reference *)
      [InlineComment (StoreArray (basetypeof dec), node2str node)]

    | _ -> raise InvalidNode
  in
  let instrs = trav program in

  (* Sort aggregated constants and add definitions
   * We might want to rewrite this 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 fst sorted_pairs in

  instrs @ const_defs

let phase = function
  | Ast node -> Assembly (assemble node)
  | _ -> raise InvalidInput