civicaml/phases/desug.ml

open Printf
open Types
open Util

(* Create new constant variables for all assigned array values so that they are
 * only evaluated once *)
let rec consts_to_vars node =
    let rec create_vars new_vars values = function
        | [] -> (new_vars, values)
        | hd :: tl ->
            let (new_vars, value) = match hd with
                | ArrayConst (values, ann) ->
                    let (new_vars, values) = create_vars new_vars [] values in
                    (new_vars, ArrayConst (values, ann))
                | value ->
                    let index = fresh_const "const" in
                    (new_vars @ [(index, value)], Var (index, None, annof value))
            in
            create_vars new_vars (values @ [value]) tl
    in
    match node with
    (* Add vardecs for values in arrayconst *)
    | VarDec (ArrayDims _ as ctype, name, Some (ArrayConst (values, vann)), ann) ->
        let (new_vars, values) = create_vars [] [] values in
        let value = ArrayConst (values, vann) in
        let create_vardec (name, value) =
            VarDec (basetypeof node, name, Some value, annof value)
        in
        let new_vardecs = List.map create_vardec new_vars in
        Block (new_vardecs @ [VarDec (ctype, name, Some value, ann)])

    (* 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 -> transform_children consts_to_vars node

(* Generate new variables for array dimensions, to avoid re-evalutation when
 * array dimensions are used (e.g., after array dimension reduction). *)
let rec array_dims node =
    let make_dims basename values make_dec =
        let make_name i _ = basename ^ "$dim$$" ^ string_of_int (i + 1) in
        let names = mapi make_name values in

        let decs = List.map2 make_dec 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
    match node with
    | VarDec (ArrayDims (ctype, values), name, init, ann) ->
        let make_dec value name = VarDec (Int, name, Some value, []) in
        let (decs, dims) = make_dims name values make_dec in
        Block (decs @ [VarDec (ArrayDims (ctype, dims), name, init, ann)])

    | GlobalDef (export, ArrayDims (ctype, values), name, None, ann) ->
        let make_dec value name = GlobalDef (export, Int, name, Some value, []) in
        let (decs, dims) = make_dims name values make_dec in
        Block (decs @ [GlobalDef (export, ArrayDims (ctype, dims), name, None, ann)])

    | GlobalDec (ArrayDims (ctype, dims), name, ann) ->
        let rec make_decs = function
            | [] -> []
            | Dim (name, ann) :: tl -> GlobalDec (Int, name, ann) :: (make_decs tl)
            | _ -> raise InvalidNode
        in
        let decs = make_decs dims in
        Block (decs @ [GlobalDec (ArrayDims (ctype, dims), name, ann)])

    | node -> transform_children array_dims node

(* Split variable declaration and initialisation *)
let rec split_inits = function
    (*
    (* Translate scalar array initialisation to ArrayScalar node, for easy
     * replacement later on *)
    | VarDec (ArrayDims (_, dims) as ctype, name, Some (Const _   as v), ann) ->
        let init = Some (ArrayInit (ArrayScalar v, dims)) in
        split_inits (VarDec (ctype, name, init, ann))

    (* Wrap ArrayConst in ArrayInit to pass dimensions *)
    | VarDec (ArrayDims (_, dims) as ctype, name, Some (ArrayConst _ as v), ann) ->
        let init = Some (ArrayInit (v, dims)) in
        split_inits (VarDec (ctype, name, init, ann))
    *)

    (* 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);
        ]

    (* Variable initialisations are split into dec;assign *)
    | 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 -> transform_children 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 -> transform_children 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
        (match extract_inits decls with
        | ([], _) -> Program (decls, ann)
        | (inits, decls) ->
            let init_func = FunDef (true, Void, "__init", [], Block inits, []) in
            Program (init_func :: decls, ann)
        )

    (* 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 -> transform_children move_inits node

let for_to_while node =
    let rec replace_var var replacement node =
        let trav = (replace_var var replacement) in
        match node with
        | Var (name, None, ann) when name = var ->
            Var (replacement, None, ann)
        | For (counter, start, stop, step, body, ann) when counter = var ->
            For (replacement, trav start, trav stop, trav step, trav body, ann)
        | node ->
            transform_children trav node
    in
    let rec traverse new_vars = function
        | FunDef (export, ret_type, name, params, body, ann) ->
            let new_vars = ref [] in
            let body = traverse 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 = new_vardecs @ (flatten_blocks (block_body body)) in
            FunDef (export, ret_type, name, params, Block _body, ann)

        (* Transform for-loops to while-loops *)
        | For (counter, start, stop, step, body, ann) ->
            let _i = fresh_var 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 0, []), []),
                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);
                traverse new_vars (While (cond, (Block (
                    block_body (replace_var counter _i body) @
                    [Assign (_i, None, Binop (Add, vi, vstep, []), [])]
                )), ann));
            ]

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

        | node -> transform_children (traverse new_vars) node
    in
    traverse (ref []) node

let rec sublist n = function
    | [] when n > 0  -> raise (Invalid_argument "n")
    | []             -> []
    | lst when n = 0 -> lst
    | _ :: tl        -> sublist (n - 1) tl

let rec array_init = function
    (* Transform array constant inisialisation into separate assign statements
     * for all entries in the constant array *)
    (* TODO: only allow when array dimensions are constant? *)
    | Assign (name, None, ArrayInit (ArrayConst _ as value, dims), ann) ->
        let ndims = List.length dims in
        let rec make_assigns depth i indices = function
            | [] -> []
            | hd :: tl ->
                let assigns = traverse depth (i :: indices) hd in
                make_assigns depth (i + 1) indices tl @ assigns
        and traverse depth indices = function
            | ArrayConst (values, _) ->
                make_assigns (depth + 1) 0 indices values
            | value when depth = ndims ->
                let indices = List.map (fun i -> Const (IntVal i, [])) indices in
                [Assign (name, Some (List.rev indices), value, ann)]
            | value when depth < ndims ->
                (* Use the for-loops constructed for scalar assignment *)
                let value = ArrayInit (value, dims) in
                let indices = List.map (fun i -> Const (IntVal i, [])) indices in
                [array_init (Assign (name, Some (List.rev indices), value, ann))]
            | node ->
                let msg = sprintf
                    "dimension mismatch: expected %d nesting levels, got %d"
                    ndims depth
                in
                raise (NodeError (node, msg))
        in
        Block (List.rev (traverse 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) as node ->
        let rec add_loop indices = function
            | [] ->
                array_init (Assign (name, Some indices, value, ann))
            | dim :: rest ->
                let counter = fresh_var "i" in
                let start = Const (IntVal 0, []) in
                let step = Const (IntVal 1, []) 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 dims_left = sublist (List.length indices) dims in
        add_loop indices dims_left

    | node -> transform_children array_init node

let phase = function
    | Ast node ->
        (* Generate variable declarations for expressions that must be evaluated
         * once and used multiple times *)
        let node = consts_to_vars (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 "desugar")