civicaml/util.ml

open Printf
open Str
open Lexing
open Types
open Globals

(** Empty location, use when node location is unkown or irrelevant. *)
let noloc = ("", 0, 0, 0, 0)

let rec repeat s n = if n < 1 then "" else s ^ (repeat s (n - 1))
let expand n text = text ^ repeat " " (n - String.length text)

(* Logging functions *)

let hline = "-----------------------------------------------------------------"

let prt_line = prerr_endline

let prt_node node = prt_line (Stringify.node2str node)

let log_plain_line verbosity line =
  if args.verbose >= verbosity then prt_line line

let log_line verbosity line =
  log_plain_line verbosity (repeat " " 13 ^ line)

let log_node verbosity node =
  if args.verbose >= verbosity then prt_node node

(* Variable generation *)
let var_counter = ref 0
let fresh_var prefix =
  var_counter := !var_counter + 1;
  prefix ^ "$" ^ string_of_int !var_counter

(* Constants are marked by a double $$ for recognition during constant
 * propagation *)
let fresh_const prefix = fresh_var (prefix ^ "$")

let loc_from_lexpos pstart pend =
  let (fname, ystart, yend, xstart, xend) = begin
    pstart.pos_fname,
    pstart.pos_lnum,
    pend.pos_lnum,
    (pstart.pos_cnum - pstart.pos_bol + 1),
    (pend.pos_cnum - pend.pos_bol)
  end in
  if ystart = yend && xend < xstart then
    (fname, ystart, yend, xstart, xstart)
  else
    (fname, ystart, yend, xstart, xend)

let rec flatten_blocks lst =
  let flatten = flatten_blocks in
  let rec trav = function
    | Block body ->
      Block (flatten body)
    | FunDef (export, ret_type, name, params, body, ann) ->
      FunDef (export, ret_type, name, flatten params, trav body, ann)
    | If (cond, body, ann) ->
      If (cond, trav body, ann)
    | IfElse (cond, tbody, fbody, ann) ->
      IfElse (cond, trav tbody, trav fbody, ann)
    | While (cond, body, ann) ->
      While (cond, trav body, ann)
    | DoWhile (cond, body, ann) ->
      DoWhile (cond, trav body, ann)
    | For (counter, start, stop, step, body, ann) ->
      For (counter, start, stop, step, trav body, ann)
    | VarDecs decs ->
      VarDecs (flatten decs)
    | LocalFuns decs ->
      LocalFuns (flatten decs)
    | node -> node
  in
  match lst with
  | []                -> []
  | Block nodes :: tl -> flatten nodes @ (flatten tl)
  | DummyNode :: tl   -> flatten tl
  | hd :: tl          -> trav hd :: (flatten tl)

(* Default tree transformation
 * (node -> node) -> node -> node *)
let transform_children trav node =
  let trav_all nodes = List.map trav nodes in
  match node with
  | Program (decls, ann) ->
    Program (flatten_blocks (trav_all decls), ann)
  | FunDec (ret_type, name, params, ann) ->
    FunDec (ret_type, name, trav_all params, ann)
  | FunDef (export, ret_type, name, params, body, ann) ->
    FunDef (export, ret_type, name, trav_all params, trav body, ann)
  | GlobalDec (ctype, name, ann) ->
    GlobalDec (ctype, name, ann)
  | GlobalDef (export, ctype, name, Some init, ann) ->
    GlobalDef (export, ctype, name, Some (trav init), ann)

  | VarDecs decs ->
    VarDecs (trav_all decs)
  | LocalFuns funs ->
    LocalFuns (trav_all funs)

  | VarDec (ctype, name, Some init, ann) ->
    VarDec (ctype, name, Some (trav init), ann)
  | Assign (name, None, value, ann) ->
    Assign (name, None, trav value, ann)
  | Assign (name, Some dims, value, ann) ->
    Assign (name, Some (trav_all dims), trav value, ann)
  | VarLet (dec, None, value, ann) ->
    VarLet (dec, None, trav value, ann)
  | VarLet (dec, Some dims, value, ann) ->
    VarLet (dec, Some (trav_all dims), trav value, ann)
  | Return (value, ann) ->
    Return (trav value, ann)
  | If (cond, body, ann) ->
    If (trav cond, trav body, ann)
  | IfElse (cond, true_body, false_body, ann) ->
    IfElse (trav cond, trav true_body, trav false_body, ann)
  | While (cond, body, ann) ->
    While (trav cond, trav body, ann)
  | DoWhile (cond, body, ann) ->
    DoWhile (trav cond, trav body, ann)
  | For (counter, start, stop, step, body, ann) ->
    For (counter, trav start, trav stop, trav step, trav body, ann)
  | Allocate (dec, dims, ann) ->
    Allocate (dec, trav_all dims, ann)
  | Expr value ->
    Expr (trav value)
  | Block (body) ->
    Block (trav_all body)

  | Monop (op, value, ann) ->
    Monop (op, trav value, ann)
  | Binop (op, left, right, ann) ->
    Binop (op, trav left, trav right, ann)
  | Cond (cond, true_expr, false_expr, ann) ->
    Cond (trav cond, trav true_expr, trav false_expr, ann)
  | TypeCast (ctype, value, ann) ->
    TypeCast (ctype, trav value, ann)
  | FunCall (name, args, ann) ->
    FunCall (name, trav_all args, ann)
  | Arg value ->
    Arg (trav value)

  | ArrayInit (value, dims) ->
    ArrayInit (trav value, dims)
  | ArrayScalar value ->
    ArrayScalar (trav value)
  | Var (dec, Some dims, ann) ->
    Var (dec, Some (trav_all dims), ann)
  | VarUse (dec, Some dims, ann) ->
    VarUse (dec, Some (trav_all dims), ann)
  | FunUse (dec, params, ann) ->
    FunUse (dec, trav_all params, ann)

  | _ -> node

let annotate a = function
  | Program (decls, ann) ->
    Program (decls, a :: ann)
  | FunDec (ret_type, name, params, ann) ->
    FunDec (ret_type, name, params, a :: ann)
  | FunDef (export, ret_type, name, params, body, ann) ->
    FunDef (export, ret_type, name, params, body, a :: ann)
  | GlobalDec (ctype, name, ann) ->
    GlobalDec (ctype, name, a :: ann)
  | GlobalDef (export, ctype, name, init, ann) ->
    GlobalDef (export, ctype, name, init, a :: ann)
  | VarDec (ctype, name, init, ann) ->
    VarDec (ctype, name, init, a :: ann)
  | Assign (name, dims, value, ann) ->
    Assign (name, dims, value, a :: ann)
  | VarLet (dec, dims, value, ann) ->
    VarLet (dec, dims, value, a :: ann)
  | Return (value, ann) ->
    Return (value, a :: ann)
  | If (cond, body, ann) ->
    If (cond, body, a :: ann)
  | IfElse (cond, true_body, false_body, ann) ->
    IfElse (cond, true_body, false_body, a :: ann)
  | While (cond, body, ann) ->
    While (cond, body, a :: ann)
  | DoWhile (cond, body, ann) ->
    DoWhile (cond, body, a :: ann)
  | For (counter, start, stop, step, body, ann) ->
    For (counter, start, stop, step, body, a :: ann)
  | Allocate (dec, dims, ann) ->
    Allocate (dec, dims, a :: ann)
  | Monop (op, value, ann) ->
    Monop (op, value, a :: ann)
  | Binop (op, left, right, ann) ->
    Binop (op, left, right, a :: ann)
  | Cond (cond, true_expr, false_expr, ann) ->
    Cond (cond, true_expr, false_expr, a :: ann)
  | TypeCast (ctype, value, ann) ->
    TypeCast (ctype, value, a :: ann)
  | FunCall (name, args, ann) ->
    FunCall (name, args, a :: ann)
  | Arg value ->
    Arg (value)
  | Var (dec, dims, ann) ->
    Var (dec, dims, a :: ann)
  | VarUse (dec, dims, ann) ->
    VarUse (dec, dims, a :: ann)
  | FunUse (dec, params, ann) ->
    FunUse (dec, params, a :: ann)
  | Const (BoolVal value, ann) ->
    Const (BoolVal value, a :: ann)
  | Const (IntVal value, ann) ->
    Const (IntVal value, a :: ann)
  | Const (FloatVal value, ann) ->
    Const (FloatVal value, a :: ann)
  | ArrayConst (value, ann) ->
    ArrayConst (value, a :: ann)
  | Param (ctype, name, ann) ->
    Param (ctype, name, a :: ann)
  | Dim (name, ann) ->
    Dim (name, a :: ann)

  | _ -> raise InvalidNode

let rec annof = function
  | Program (_, ann)
  | Param (_, _, ann)
  | Dim (_, ann)
  | FunDec (_, _, _, ann)
  | FunDef (_, _, _, _, _, ann)
  | GlobalDec (_, _, ann)
  | GlobalDef (_, _, _, _, ann)
  | VarDec (_, _, _, ann)
  | Assign (_, _, _, ann)
  | VarLet (_, _, _, ann)
  | Return (_, ann)
  | If (_, _, ann)
  | IfElse (_, _, _, ann)
  | While (_, _, ann)
  | DoWhile (_, _, ann)
  | For (_, _, _, _, _, ann)
  | Allocate (_, _, ann)
  | Const (BoolVal _, ann)
  | Const (IntVal _, ann)
  | Const (FloatVal _, ann)
  | ArrayConst (_, ann)
  | Var (_, _, ann)
  | Monop (_, _, ann)
  | Binop (_, _, _, ann)
  | Cond (_, _, _, ann)
  | TypeCast (_, _, ann)
  | VarUse (_, _, ann)
  | FunUse (_, _, ann)
  | FunCall (_, _, ann) -> ann

  | ArrayInit (value, _)
  | ArrayScalar value
  | Expr value
  | Arg value -> annof value

  | _ -> raise InvalidNode

let locof node =
  let rec trav = function
    | []           -> noloc
    | Loc loc :: _ -> loc
    | _ :: tl      -> trav tl
  in trav (annof node)

let depthof node =
  let rec trav = function
    | [] ->
      prerr_string "cannot get depth for: ";
      prt_node node;
      raise InvalidNode
    | Depth depth :: _ -> depth
    | _ :: tl          -> trav tl
  in trav (annof node)

let indexof node =
  let rec trav = function
    | [] ->
      prerr_string "cannot get index for: ";
      prt_node node;
      raise InvalidNode
    | Index index :: _ -> index
    | _ :: tl          -> trav tl
  in trav (annof node)

let typeof = function
  (* Some nodes have their type as property *)
  | VarDec (ctype, _, _, _)
  | Param (ctype, _, _)
  | FunDec (ctype, _, _, _)
  | FunDef (_, ctype, _, _, _, _)
  | GlobalDec (ctype, _, _)
  | GlobalDef (_, ctype, _, _, _)
  | TypeCast (ctype, _, _)
    -> ctype

  (* Dim nodes are always type Int, and are copied by context analysis before
   * they are annotated with Type Int, so this match is necessary *)
  | Dim _ -> Int

  (* Other nodes must be annotated during typechecking *)
  | node ->
    let rec trav = function
      | [] ->
        prerr_string "cannot get type for: ";
        prt_node node;
        raise InvalidNode
      | Type t :: _ -> t
      | _ :: tl     -> trav tl
    in trav (annof node)

let labelof node =
  let rec trav = function
    | [] ->
      prerr_string "cannot get label for: ";
      prt_node node;
      raise InvalidNode
    | LabelName label :: _ -> label
    | _ :: tl              -> trav tl
  in trav (annof node)

let const_type = function
  | BoolVal _  -> Bool
  | IntVal _   -> Int
  | FloatVal _ -> Float

(*
let get_line str n =
  let rec find_start from = function
    | n when n < 1 -> raise (Invalid_argument "n")
    | 1 -> from
    | n -> find_start ((String.index_from str from '\n') + 1) (n - 1)
  in
  let linestart = find_start 0 n in
  let len = String.length str in
  let lineend =
    try String.index_from str linestart '\n'
    with Not_found -> len
  in
  String.sub str linestart (lineend - linestart)
*)

let count_tabs str upto =
  let rec count n = function
    | 0 -> n
    | i -> count (if String.get str (i - 1) = '\t' then n + 1 else n) (i - 1)
  in count 0 upto

let tabwidth = 4

let retab str = global_replace (regexp "\t") (repeat " " tabwidth) str

let indent n = repeat (repeat " " (tabwidth - 1)) n

let prerr_loc (fname, ystart, yend, xstart, xend) =
  let file = open_in fname in

  (* skip lines until the first matched line *)
  for i = 1 to ystart - 1 do let _ = input_line file in () done;

  (* for each line in `loc`, print the source line with an underline *)
  for l = ystart to yend do
    let line = input_line file in
    let linewidth = String.length line in
    let left = if l = ystart then xstart else 1 in
    let right = if l = yend then xend else linewidth in
    if linewidth > 0 then begin
      prerr_endline (retab line);
      prerr_string (indent (count_tabs line right));
      for i = 1 to left - 1 do prerr_char ' ' done;
      for i = left to right do prerr_char '^' done;
      prerr_endline "";
    end
  done;
  ()

let prerr_loc_msg loc msg =
  if args.verbose >= 1 then begin
    let (fname, ystart, yend, xstart, xend) = loc in
    if loc != noloc then begin
      let line_s = if yend != ystart
        then sprintf "lines %d-%d" ystart yend
        else sprintf "line %d" ystart
      in
      let char_s = if xend != xstart || yend != ystart
        then sprintf "characters %d-%d" xstart xend
        else sprintf "character %d" xstart
      in
      eprintf "File \"%s\", %s, %s:\n" fname line_s char_s;
    end;
    eprintf "%s\n" msg;

    if args.verbose >= 1 && loc != noloc then
        try prerr_loc loc
        with Sys_error _ -> ()
  end;
  ()

let block_body = function
  | Block nodes -> nodes
  | _ -> raise InvalidNode

let basetypeof node = match typeof node with
  | ArrayDims (ctype, _)
  | Array ctype
  | ctype -> ctype

let nameof = function
  | GlobalDec (_, name, _)
  | GlobalDef (_, _, name, _, _)
  | FunDec (_, name, _, _)
  | FunDef (_, _, name, _, _, _)
  | VarDec (_, name, _, _)
  | Param (_, name, _)
  | Dim (name, _) -> name
  | _ -> raise InvalidNode

let optmap f = function
  | None -> None
  | Some lst -> Some (List.map f lst)

let optmapl f = function
  | None -> []
  | Some lst -> List.map f lst

let mapi f lst =
  let rec trav i = function
    | [] -> []
    | hd :: tl -> f i hd :: (trav (i + 1) tl)
  in trav 0 lst

(** Constants that are *)
let immediate_consts = [
  BoolVal true;
  BoolVal false;
  IntVal (-1);
  IntVal 0;
  IntVal 1;
  FloatVal 0.0;
  FloatVal 1.0;
]

let is_immediate_const const =
  if args.optimize then List.mem const immediate_consts else false

let is_array node = match typeof node with
  | ArrayDims _ | Array _ -> true
  | _ -> false

let node_warning node msg =
  prerr_loc_msg (locof node) ("Warning: " ^ msg)