civicaml/phases/desug.ml

open Printf
open Types
open Util

(* Check if a function defines a variable name *)
let defines var = function
  | FunDef (export, ret_type, name, params, Block (VarDecs decs :: tl), ann) ->
    let rec trav_decs = function
      | [] -> false
      | Param (ArrayDims (_, dims), name, _) :: tl ->
        name = var || trav_decs dims || trav_decs tl
      | (Dim (name, _) | VarDec (_, name, _, _) | Param (_, name, _)) :: _
        when name = var -> true
      | _ :: tl -> trav_decs tl
    in
    trav_decs params || trav_decs decs
  | _ -> raise InvalidNode

(* Replace all occurences of a variable name with another name *)
let rec replace_var var replacement node =
  let trav = (replace_var var replacement) in
  let trav_all = List.map trav in
  let trav_opt = function None -> None | Some node -> Some (trav node) in
  match node with
  (* Replace variable name on match *)
  | Var (name, ind, ann) when name = var ->
    let ind = match ind with None -> None | Some ind -> Some (trav_all ind) in
    Var (replacement, ind, ann)

  (* Don't enter a function body if it redefines the variable *)
  | FunDef _ when defines var node -> node

  (* Don't traverse into a for-loop body if the loop counter redefines var *)
  | For (counter, start, stop, step, body, ann) when counter = var ->
    For (counter, trav start, trav stop, trav step, body, ann)

  (* At this point, array dimension expressions may not have been moved to new
   * variables yet, so traverse them explicitly *)
  | VarDec (ArrayDims (ctype, dims), name, init, ann) ->
    VarDec (ArrayDims (ctype, trav_all dims), name, trav_opt init, ann)

  | node -> traverse_unit trav node

(* Create new constant variables for scalar initialisations on arrays so that
 * they are only evaluated once *)
let rec move_scalars = function
  (* Prevent next match for ArrayConst initialisations *)
  | VarDec (ArrayDims _, _, Some (ArrayConst _), _) as node ->
    node

  (* Add vardec for scalar value *)
  | VarDec (ArrayDims _ as ctype, name, Some value, ann) as node ->
    let scalar_name = fresh_const "scalar" in
    Block [
      VarDec (basetypeof node, scalar_name, Some value, ann);
      VarDec (ctype, name, Some (Var (scalar_name, None, annof value)), ann);
    ]

  | node -> traverse_unit move_scalars node

(* Generate new variables for array dimensions, to avoid re-evalutation when
 * array dimensions are used (e.g., after array dimension reduction). *)
let array_dims node =
  (*
  let make_dimname basename i = generate_const basename (i + 1) in
  *)
  let make_dimname = generate_const in
  let patch_dims basename values make_decs =
    let names = mapi (fun i _ -> make_dimname basename i) values in

    let decs = List.concat (List.map2 make_decs values names) in

    let make_dim value name = Dim (name, annof value) in
    let dims = List.map2 make_dim values names in

    (decs, dims)
  in

  (* Save dimension replacements in one global hash table (we are not replacing
   * local vars, so everything is in the global scope) *)
  let replacements = Hashtbl.create 10 in

  let rec trav = function
    | VarDec (ArrayDims (ctype, values), name, init, ann) ->
      let make_decs value name = [VarDec (Int, name, Some value, [])] in
      let decs, dims = patch_dims name values make_decs in
      Block (decs @ [VarDec (ArrayDims (ctype, dims), name, init, ann)])

    | GlobalDef (export, ArrayDims (ctype, values), name, init, ann) ->
      (* Move array dimensions into new variables to avoid double evaluation of
      * expressions with side effects (i.e. function calls) *)
      let make_decs value name = [GlobalDef (export, Int, name, Some value, [])] in
      let decs, dims = patch_dims name values make_decs in
      Block (decs @ [GlobalDef (export, ArrayDims (ctype, dims), name, init, ann)])

    | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
      (* Create an 'extern int ...' definition for each dimension with a
      * consistent import name, and replace all local uses with the imported
      * variable name *)
      let make_decs dim impname =
        match dim with
        | Dim (dimname, _) ->
          (* Fix name clashes (needed because context analysis has not been done
           * yet) *)
          if Hashtbl.mem replacements dimname then begin
            raise (FatalError (NodeMsg (dim, "duplicate dimension name")))
          end;

          (* Occurences of dimension names are replaced after the traversal *)
          Hashtbl.add replacements dimname impname;

          [GlobalDec (Int, impname, [])]
        | _ -> raise InvalidNode
      in
      let decs, dims = patch_dims name dims make_decs in
      Block (decs @ [GlobalDec (ArrayDims (ctype, dims), name, ann)])

      (*
      let make_decs i = function
        | Dim (dimname, dimann) ->
          let impname = generate_id name (i + 1) in
          let decs = [
            GlobalDec (Int, impname, []);
            GlobalDef (false, Int, dimname, Some (Var (impname, None, [])), [])
          ] in
          (decs, Dim (impname, dimann))
        | _ -> raise InvalidNode
      in
      let decs, dims = List.split (mapi make_decs dims) in
      Block (List.concat decs @ [GlobalDec (ArrayDims (ctype, dims), name, ann)])
      *)

    | node -> traverse_unit trav node
  in
  Hashtbl.fold replace_var replacements (trav node)

(* Split variable initialisation into declaration and assignment *)
let rec split_inits = function
  (* Wrap array initialisation in ArrayInit to pass dimensions *)
  | VarDec (ArrayDims (_, dims) as ctype, name, Some value, ann) ->
    Block [
      VarDec (ctype, name, None, ann);
      Assign (name, None, ArrayInit (value, dims), ann);
    ]

  | GlobalDef (export, (ArrayDims (_, dims) as ctype), name, Some value, ann) ->
    Block [
      GlobalDef (export, ctype, name, None, ann);
      Assign (name, None, ArrayInit (value, dims), ann);
    ]

  | VarDec (ctype, name, Some init, ann) ->
    Block [
      VarDec (ctype, name, None, ann);
      Assign (name, None, init, ann);
    ]

  | GlobalDef (export, ctype, name, Some init, ann) ->
    Block [
      GlobalDef (export, ctype, name, None, ann);
      Assign (name, None, init, ann);
    ]

  | node -> traverse_unit split_inits node

(* Add <allocate> statements after array declarations *)
let rec add_allocs node =
  let create_dimvar = function
    | Dim (name, _) -> Var (name, None, [])
    | _ -> raise InvalidNode
  in
  match node with
  | VarDec (ArrayDims (_, dims), _, _, ann) ->
    Block [node; Allocate (node, List.map create_dimvar dims, ann)]

  | GlobalDef (_, ArrayDims (_, dims), _, _, ann) ->
    Block [node; Allocate (node, List.map create_dimvar dims, ann)]

  | node -> traverse_unit add_allocs node

let extract_inits lst =
  let rec trav inits = function
    | [] ->
      (List.rev inits, [])
    | (Assign _ as hd) :: tl
    | (Allocate _ as hd) :: tl ->
      trav (hd :: inits) tl
    | hd :: tl ->
      let inits, tl = trav inits tl in
      (inits, (hd :: tl))
  in trav [] lst

let rec move_inits = function
  (* Move global initialisations to __init function *)
  | Program (decls, ann) ->
    let decls = List.map move_inits decls in
    begin match extract_inits decls with
    | ([], _) -> Program (decls, ann)
    | (inits, decls) ->
      let body = Block (VarDecs [] :: LocalFuns [] :: inits) in
      let init_func = FunDef (true, Void, "__init", [], body, []) in
      Program (init_func :: decls, ann)
    end

  (* Split local variable initialisations in declaration and assignment *)
  | FunDef (export, ret_type, name, params, Block body, ann) ->
    let rec place_inits inits = function
      | VarDecs lst :: tl ->
        let inits, decs = extract_inits lst in
        VarDecs decs :: (place_inits inits tl)
      | LocalFuns _ as hd :: tl ->
        hd :: inits @ tl
      | _ -> raise InvalidNode
    in
    let body = Block (place_inits [] body) in
    FunDef (export, ret_type, name, params, body, ann)

  | node -> traverse_unit move_inits node

let for_to_while node =
  let rec trav new_vars = function
    | FunDef (export, ret_type, name, params, body, ann) ->
      let rec place_decs decs = function
        | Block (VarDecs lst :: tl) -> Block (VarDecs (decs @ lst) :: tl)
        | _ -> raise InvalidNode
      in
      let new_vars = ref [] in
      let body = trav new_vars body in
      let create_vardec name = VarDec (Int, name, None, []) in
      let new_vardecs = List.map create_vardec !new_vars in
      let body = place_decs new_vardecs body in
      FunDef (export, ret_type, name, params, body, ann)

    (* Transform for-loops to while-loops *)
    | For (counter, start, stop, step, body, ann) ->
      let _i = fresh_id counter in
      let _stop = fresh_const "stop" in
      let _step = fresh_const "step" in
      new_vars := !new_vars @ [_i; _stop; _step];

      let vi = Var (_i, None, []) in
      let vstop = Var (_stop, None, annof stop) in
      let vstep = Var (_step, None, annof step) in
      let cond = Cond (
        Binop (Gt, vstep, Const (IntVal 0l, []), []),
        Binop (Lt, vi, vstop, []),
        Binop (Gt, vi, vstop, []),
        []
      ) in
      Block [
        Assign (_i, None, start, annof start);
        Assign (_stop, None, stop, annof stop);
        Assign (_step, None, step, annof step);
        trav new_vars (While (cond, (Block (
          block_body (replace_var counter _i body) @
          [Assign (_i, None, Binop (Add, vi, vstep, []), [])]
        )), ann));
      ]

    (* Transform while-loops to do-while loops in if-statements *)
    (* DISABLED, while-loops are explicitly supported by the assembly phase
    | While (cond, body, ann) ->
      let cond = trav new_vars cond in
      let body = trav new_vars body in
      Block [If (cond, Block [DoWhile (cond, body, ann)], ann)]
    *)

    | node -> traverse_unit (trav new_vars) node
  in
  trav (ref []) node

let rec array_init = function
  (* Transform array constant initialisation into separate assign statements
   * for all entries in the array literal *)
  | Assign (name, None, ArrayInit (ArrayConst _ as value, dims), ann) ->
    let intconst i = Const (IntVal (Int32.of_int i), []) in
    let ndims = List.length dims in
    let rec make_assigns depth i indices = function
      | [] -> []
      | hd :: tl ->
        let assigns = trav depth (i :: indices) hd in
        make_assigns depth (i + 1) indices tl @ assigns
    and trav depth indices = function
      | ArrayConst (values, _) ->
        make_assigns (depth + 1) 0 indices values
      | value when depth = ndims ->
        let indices = List.map intconst indices in
        [Assign (name, Some (List.rev indices), value, ann)]
      (* DISABLED: nesting level must now be equal to number of dimensions
      | value when depth < ndims ->
        (* Use for-loops for scalar assignment on sub-array *)
        let value = ArrayInit (value, dims) in
        let indices = List.map intconst indices in
        [array_init (Assign (name, Some (List.rev indices), value, ann))]
      *)
      | node ->
        raise (FatalError (NodeMsg (node, sprintf
          "dimension mismatch: expected %d nesting levels, got %d"
          ndims depth)))
    in
    Block (List.rev (trav 0 [] value))

  (* Replace no indices with empty indices to have a list below *)
  | Assign (name, None, (ArrayInit _ as value), ann) ->
    array_init (Assign (name, Some [], value, ann))

  | Assign (name, Some indices, ArrayInit (value, dims), ann) ->
    let rec add_loop indices = function
      | [] ->
        array_init (Assign (name, Some indices, value, ann))
      | dim :: rest ->
        let counter = fresh_id "i" in
        let start = Const (IntVal 0l, []) in
        let step = Const (IntVal 1l, []) in
        let body = Block [add_loop (indices @ [Var (counter, None, [])]) rest] in
        let stop = match dim with
        | Dim (name, ann) -> Var (name, None, ann)
        | _ -> dim
        in
        For (counter, start, stop, step, body, [])
    in
    let rec sublist n = function
      | [] when n > 0  -> raise (Invalid_argument "n")
      | []             -> []
      | lst when n = 0 -> lst
      | _ :: tl        -> sublist (n - 1) tl
    in
    let dims_left = sublist (List.length indices) dims in
    add_loop indices dims_left

  | node -> traverse_unit array_init node

let phase = function
  | Ast node ->
    (* Move array dimensions and scalar initialisations into new variables as
     * initialisations, so that they are evaluated exactly once, and so that
     * dimension names are consistent with the array name *)
    let node = move_scalars (array_dims node) in

    (* Split variable initialisations into declarations and assignments, and
     * move the assignments to the function body *)
    let node = move_inits (add_allocs (split_inits node)) in

    (* Transform ArrayConst assignment to assignments in for-loops, and
     * transform all for-loops to while-loops afterwards *)
    Ast (for_to_while (array_init node))

  | _ -> raise InvalidInput