open Lexing
open Ast

let var_counter = ref 0

let fresh_var prefix =
    var_counter := !var_counter + 1;
    prefix ^ "$" ^ string_of_int !var_counter

let loc_from_lexpos pstart pend =
    let (fname, ystart, yend, xstart, xend) = (
        pstart.pos_fname,
        pstart.pos_lnum,
        pend.pos_lnum,
        (pstart.pos_cnum - pstart.pos_bol + 1),
        (pend.pos_cnum - pend.pos_bol)
    ) in
    if ystart = yend && xend < xstart then
        (fname, ystart, yend, xstart, xstart)
    else
        (fname, ystart, yend, xstart, xend)

(* Default tree transformation
 * (node -> node) -> node -> node *)
let transform_children trav node =
    let trav_all nodes = List.map trav nodes in
    match node with
    | Program (decls, loc) ->
        Program (trav_all decls, loc)
    | FunDec (ret_type, name, params, loc) ->
        FunDec (ret_type, name, trav_all params, loc)
    | FunDef (export, ret_type, name, params, body, loc) ->
        FunDef (export, ret_type, name, trav_all params, trav body, loc)
    | GlobalDec (ctype, name, loc) ->
        GlobalDec (ctype, name, loc)
    | GlobalDef (export, ctype, name, Some init, loc) ->
        GlobalDef (export, ctype, name, Some (trav init), loc)

    | VarDec (ctype, name, Some init, loc) ->
        VarDec (ctype, name, Some (trav init), loc)
    | Assign (name, value, loc) ->
        Assign (name, trav value, loc)
    | Return (value, loc) ->
        Return (trav value, loc)
    | If (cond, body, loc) ->
        If (trav cond, trav body, loc)
    | IfElse (cond, true_body, false_body, loc) ->
        IfElse (trav cond, trav true_body, trav false_body, loc)
    | While (cond, body, loc) ->
        While (trav cond, trav body, loc)
    | DoWhile (cond, body, loc) ->
        DoWhile (trav cond, trav body, loc)
    | For (counter, start, stop, step, body, loc) ->
        For (counter, trav start, trav stop, trav step, trav body, loc)
    | Expr value ->
        Expr (trav value)

    | Monop (op, value, loc) ->
        Monop (op, trav value, loc)
    | Binop (op, left, right, loc) ->
        Binop (op, trav left, trav right, loc)
    | Cond (cond, true_expr, false_expr, loc) ->
        Cond (trav cond, trav true_expr, trav false_expr, loc)
    | TypeCast (ctype, value, loc) ->
        TypeCast (ctype, trav value, loc)
    | FunCall (name, args, loc) ->
        FunCall (name, trav_all args, loc)

    | Block (body) ->
        Block (trav_all body)

    | VarUse (var, def, depth) ->
        VarUse (trav var, def, depth)

    | _ -> node

 let rec locof = function
    | Program (_, loc)
    | Param (_, _, loc)
    | FunDec (_, _, _, loc)
    | FunDef (_, _, _, _, _, loc)
    | GlobalDec (_, _, loc)
    | GlobalDef (_, _, _, _, loc)
    | VarDec (_, _, _, loc)
    | Assign (_, _, loc)
    | Return (_, loc)
    | If (_, _, loc)
    | IfElse (_, _, _, loc)
    | While (_, _, loc)
    | DoWhile (_, _, loc)
    | For (_, _, _, _, _, loc)
    | Allocate (_, _, loc)
    | BoolConst (_, loc)
    | IntConst (_, loc)
    | FloatConst (_, loc)
    | ArrayConst (_, loc)
    | ArrayScalar (_, loc)
    | Var (_, loc)
    | Deref (_, _, loc)
    | Monop (_, _, loc)
    | Binop (_, _, _, loc)
    | Cond (_, _, _, loc)
    | TypeCast (_, _, loc)
    | FunCall (_, _, loc)
        -> loc

    | Expr value
    | VarUse (value, _, _)
        -> locof value

    | Block _ -> noloc