open Printf
open Types
open Util

let rec var_init = function
    (* Move global initialisations to __init function *)
    | Program (decls, ann) ->
        let decls = flatten_blocks (List.map var_init decls) in
        let rec trav assigns = function
            | [] -> (assigns, [])
            | (Assign _ as hd) :: tl
            | (Allocate _ as hd) :: tl -> trav (assigns @ [hd]) tl
            | hd :: tl ->
                let (assigns, decls) = trav assigns tl in
                (assigns, (hd :: decls))
        in
        let (assigns, decls) = trav [] decls in (
            match assigns with
            | [] -> Program (decls, ann)
            | assigns ->
                let init_func = FunDef (true, Void, "__init", [], Block assigns, []) in
                Program (init_func :: decls, ann)
        )

    (* Global variable initialisation:
     * Add an assign statement and the Program node will remove it later on *)
    | GlobalDef (export, ctype, name, Some init, ann) ->
        Block [GlobalDef (export, ctype, name, None, ann);
               Assign (name, None, init, ann)]

    (* Global array definition:
     * - Create a new global variable for each dimension and initialise it to
     *   the given expression
     * - create __allocate statement in __init *)
    | GlobalDef (export, Array (ctype, dims), name, None, ann) as dec ->
        let rec create_dimvars i = function
            | [] -> []
            | hd :: tl ->
                let dimname = name ^ "$" ^ string_of_int i in
                let var = Var (dimname, None, ann) in
                var :: (create_dimvars (i + 1) tl)
        in
        let dimvars = create_dimvars 1 dims in
        let create_globaldef dim = function
            | Var (dimname, None, ann) ->
                var_init (GlobalDef (export, Int, dimname, Some dim, ann))
            | _ -> raise InvalidNode
        in
        let vardecs = List.map2 create_globaldef dims dimvars in
        let alloc = [Allocate (dec, dimvars, ann)] in
        Block (vardecs @
               [GlobalDef (export, Array (ctype, dimvars), name, None, ann)] @
               alloc)

    (* Split local variable initialisations in declaration and assignment *)
    | FunDef (export, ret_type, name, params, body, ann) ->
        let inits = ref [] in
        let rec extract_inits = function
            (* Translate scalar array initialisation to ArrayScalar node,
                * for easy replacement later on *)
            | VarDec (Array _ as vtype, name, Some (Const _   as v), ann) ->
                let init = Some (ArrayInit (ArrayScalar v, vtype)) in
                extract_inits (VarDec (vtype, name, init, ann))

            (* Wrap ArrayConst in ArrayInit to pass dimensions *)
            | VarDec (Array _ as vtype, name, Some (ArrayConst _ as v), ann) ->
                let init = Some (ArrayInit (v, vtype)) in
                extract_inits (VarDec (vtype, name, init, ann))

            | VarDec (ctype, name, init, ann) as dec ->
                (* array definition: create __allocate statement *)
                let alloc = match ctype with
                    | Array (_, dims) -> [Allocate (dec, dims, ann)]
                    | _ -> []
                in
                (* initialisation: create assign statement *)
                let add = match init with
                    | Some value -> alloc @ [Assign (name, None, value, ann)]
                    | None -> alloc
                in
                inits := !inits @ add;
                VarDec (ctype, name, None, ann)

            | LocalFuns funs -> LocalFuns (List.map var_init funs)

            | node -> transform_children extract_inits node
        in
        let rec place_inits = function
            (* initialisations need to be placed after local functions *)
            | (LocalFuns _ as hd) :: tl -> hd :: !inits @ tl
            | hd :: tl -> hd :: (place_inits tl)
            | [] -> []
        in
        let params = flatten_blocks (List.map var_init params) in
        let body = flatten_blocks (place_inits (block_body (extract_inits body))) in
        FunDef (export, ret_type, name, params, Block body, ann)

    | node -> transform_children var_init 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

let for_to_while node =
    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, Array (_, 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, Array (_, dims)), ann) ->
        let ndims = list_size 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

(* 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 (note the $$ which will help
 * later during constant propagation):
 * void foo() {
 *    int a$dim$$1 = 10;
 *    int a$dim$$2 = dim();
 *    int[a$dim$$1, a$dim$$2] arr;
 *    arr[1, 2] = 1;
 * }
 *
 * ... which then becomes:
 * void foo() {
 *    int a$dim$$1;
 *    int a$dim$$2;
 *    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 = function
    | VarDec (Array (ctype, dims), name, init, ann) ->
        let make_dimname i _ = name ^ "$dim$$" ^ string_of_int (i + 1) in
        let dimnames = mapi make_dimname dims in

        let make_dimvar d n = Var (n, None, annof d) in
        let dimvars = List.map2 make_dimvar dims dimnames in

        let make_dimdec dimname dim = VarDec (Int, dimname, Some dim, []) in
        let dimdecs = List.map2 make_dimdec dimnames dims in

        Block (dimdecs @ [VarDec (Array (ctype, dimvars), name, init, ann)])

    | node -> transform_children array_dims node

let rec phase input =
    log_line 1 "- Desugaring";
    match input with
    | Ast node -> Ast (for_to_while (array_init (var_init (array_dims node))))
    | _ -> raise (InvalidInput "desugar")