open Printf
open Types
open Util

(* Generate new variables for array dimensions in function bodies, to avoid
 * re-evalutation after array dimension reduction. For example:
 *
 * int dims = 0;
 *
 * int dim() {
 *     dims = dims 1;  // Side effect => dims() should be called once
 *     return 10;
 * }
 *
 * void foo() {
 *    int[10, dim()] arr;
 *    arr[0, 1] = 1;
 * }
 *
 * After dimension reduction, this would become:
 * void foo() {
 *    int[] arr;
 *    arr = allocate(10, dim());
 *    arr[1 * dim() + 0] = 1;
 * }
 *
 * This behaviour is of course incorrect. To avoid dim() from being evaluated
 * twice, the snippet above is transformed into the following code: (note the $$
 * which will help later during constant propagation)
 * void foo() {
 *    int[a$dim$$1, a$dim$$2] arr;
 *    a$dim$$1 = 10;
 *    a$dim$$2 = dim();
 *    arr[1, 2] = 1;
 * }
 *
 * ... which later becomes:
 * void foo() {
 *    int[a$dim$$1, a$dim$$2] arr;
 *    a$dim$$1 = 10;
 *    a$dim$$2 = dim();
 *    arr = __allocate(a$dim$$1 * a$dim$$2);
 *    arr[1 * a$dim$2 * 0] = 1;
 * }
 * *)
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))

    (* 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 explicittly 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 array_init = function
    (* Transform scalar assignment into nested for-loops *)
    | Assign (name, None, ArrayInit (ArrayScalar value, dims), ann) ->
        let rec add_loop indices = function
            | [] ->
                Assign (name, Some indices, value, ann)
            | dim :: rest ->
                let counter = fresh_var "i" in
                let body = Block [add_loop (indices @ [Var (counter, None, [])]) rest] in
                For (counter, Const (IntVal 0, []), dim, Const (IntVal 1, []), body, [])
        in
        add_loop [] dims

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

    | node -> transform_children array_init node

let phase = function
    | Ast node ->
        let node = move_inits (add_allocs (split_inits (array_dims node))) in
        Ast (for_to_while (array_init (node)))
    | _ -> raise (InvalidInput "desugar")