civicaml/phases/desug.ml

open Ast
open Util

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

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

                | VarDec (ctype, name, init) :: t ->
                    (* array definition: create __allocate statement *)
                    let alloc = match ctype with
                        | ArrayDef (_, dims) -> [Allocate (name, dims)]
                        | _ -> []
                    in
                    (* variable initialisation: create assign statement *)
                    let stats = match init with
                        | Some value -> alloc @ [Assign (name, value)]
                        | None -> alloc
                    in
                    VarDec (ctype, name, None) :: (trav (inits @ stats) 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, move_inits body)

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

    | Program (decls, l) ->
        let decls = flatten (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, l)
            | assigns ->
                let init_func = FunDef (true, Void, "__init", [], assigns) in
                Program (init_func :: decls, l)
            )

    | node -> transform 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)) -> Statements []

    | node -> transform array_init node
*)

let rec phase repr =
    let _ = print_endline "- Var init" in
    match repr with
        | Node (node, verbose) ->
            Node (var_init node, verbose)
        | _ -> failwith "invalid input for this phase"