open Ast
open Util

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

    (* Move global variable initialisations to exported __init function *)
    | GlobalDef (export, ctype, name, Some init, loc) ->
        Block [GlobalDef (export, ctype, name, None, loc);
               Assign (name, None, init, locof init)]

    (* Split local variable initialisations in declaration and assignment *)
    | FunDef (export, ret_type, name, params, Block body, loc) ->
        let move_inits body =
            let rec trav inits node = match node with
                (* translate scalar array initialisation to ArrayScalar node,
                 * for easy replacement later on *)
                | VarDec (ArrayDef _ as vtype, name,
                          Some (BoolConst _  as v), loc) :: t
                | VarDec (ArrayDef _ as vtype, name,
                          Some (FloatConst _ as v), loc) :: t
                | VarDec (ArrayDef _ as vtype, name,
                          Some (IntConst _   as v), loc) :: t ->
                    trav inits (VarDec (vtype, name, Some (ArrayScalar (v, vtype)), loc) :: t)

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

                (* initialisations need to be placed after local functions *)
                | (FunDef (_, _, _, _, _, _) as h) :: t ->
                    (var_init h) :: (trav inits t)

                (* rest of function body: recurse *)
                | rest -> inits @ (List.map var_init rest)
            in
            flatten_blocks (trav [] body)
        in
        let params = flatten_blocks (List.map var_init params) in
        FunDef (export, ret_type, name, params, Block (move_inits body), loc)

    | 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, loc) when name = var ->
        Var (replacement, loc)
    | For (counter, start, stop, step, body, loc) when counter = var ->
        For (replacement, trav start, trav stop, trav step, trav body, loc)
    | node ->
        transform_children trav node

let for_to_while node =
    let rec traverse new_vars = function
        | FunDef (export, ret_type, name, params, body, loc) ->
            let new_vars = ref [] in
            let body = traverse new_vars body in
            let create_vardec name = VarDec (Int, name, None, noloc) 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, loc)

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

            let vi = Var (_i, noloc) in
            let vstop = Var (_stop, locof stop) in
            let vstep = Var (_step, locof step) in
            let cond = Cond (
                Binop (Gt, vstep, IntConst (0, noloc), noloc),
                Binop (Lt, vi, vstop, noloc),
                Binop (Gt, vi, vstop, noloc),
                noloc
            ) in
            Block [
                Assign (_i, None, start, locof start);
                Assign (_stop, None, stop, locof stop);
                Assign (_step, None, step, locof step);
                While (cond, traverse new_vars (Block (
                    (* TODO: check for illegal assigments of counter in body *)
                    block_body (replace_var counter _i body) @
                    [Assign (_i, None, Binop (Add, vi, vstep, noloc), noloc)]
                )), loc);
            ]

        | 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, ArrayScalar (value, ArrayDef (_, dims)), loc) ->
        let rec add_loop indices = function
            | [] ->
                Assign (name, Some indices, value, loc)
            | dim :: rest ->
                let counter = fresh_var "i" in
                let body = Block [add_loop (indices @ [Var (counter, noloc)]) rest] in
                For (counter, IntConst (0, noloc), dim, IntConst (1, noloc), body, noloc)
        in
        add_loop [] dims

    (* TODO *)
    | Assign (name, None, ArrayConst (dims, _), _) ->
        Block []

    | node -> transform_children array_init node

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