open Ast
open Util

let rec flatten_blocks = function
    | [] -> []
    | Block nodes :: t -> (flatten_blocks nodes) @ (flatten_blocks t)
    | h :: t -> h :: (flatten_blocks t)

let rec var_init = function
    (* 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  (_, l)) as v), loc) :: t
                | VarDec (ArrayDef (_, _) as vtype, name,
                          Some ((FloatConst (_, l)) as v), loc) :: t
                | VarDec (ArrayDef (_, _) as vtype, name,
                          Some ((IntConst   (_, l)) as v), loc) :: t ->
                    trav inits (VarDec (vtype, name, Some (ArrayScalar (v, l)), loc) :: t)

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

                (* 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 trav [] body
        in
        FunDef (export, ret_type, name, params, Block (move_inits body), 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, init, locof init)]

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

    | node -> transform_children var_init node

(*
let rec array_init = function
    (* transform scalar assignment into nested for loops *)
    | Assign (name, ArrayScalar (value)) ->
        let rec add_loop indices = function
            | [] ->
                Assign (Deref (name, indices), value)
            | dim :: rest ->
                let counter = fresh_var "counter" in
                let ind = (indices @ [Var counter]) in
                For (counter, IntConst 0, dim, IntConst 1, add_loop ind rest)
        in
        add_loop [] dims

    | Assign (name, ArrayConst (dims)) -> Block []

    | node -> transform array_init node
*)

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