open Format
open Common.Pp

type ty_int = Int1 | Int32 | Int64
type ty = TInt of ty_int | TPtr | TVoid | TArray of ty * int | TDouble
type atom = Int of ty_int * int64 | Double of float | Var of string
type prim1 = Not | Neg | FNeg | Float_of_int | Int_of_float

type prim2 =
  | Eq
  | Lt
  | Le
  | Gt
  | Ge
  | And
  | Or
  | Plus
  | Minus
  | Times
  | Div
  | Mod
  | Shl
  | Lshr
  | Ashr
  | FPlus
  | FMinus
  | FTimes
  | FDiv
  | FLt
  | FLe
  | FEq
  | FGe
  | FGt

type expr =
  | Int of ty_int * int64
  | Double of float
  | Var of string
  | Let of string * expr * expr
  | Prim1 of prim1 * expr
  | Prim2 of prim2 * expr * expr
  | If of expr * expr * expr
  | Switch of expr * (int64 * expr) list * expr
  | For of string * expr * expr * expr
  | Apply of expr * (expr * ty) list * ty
  | ApplyExt of string * ty list * ty * expr list
  | Load of ty * expr
  | Store of expr * expr
  | Gep of ty * expr * int
  | Malloc of ty
  | Alloca of ty
  | Seq of expr list
  | Func of string
  | Extractelement of expr * expr
  | Insertelement of expr * expr * expr
  | Inttoptr of expr * ty
  | Ptrtoint of expr * ty
  | Intcast of expr * ty
  | Bitcast of expr * ty

type fundef = {
  name : string;
  args : (string * ty) list;
  return : ty;
  body : expr;
}

type compunit = fundef list

let pp_ty_int fmt = function
  | Int1 -> fprintf fmt "i1"
  | Int32 -> fprintf fmt "i32"
  | Int64 -> fprintf fmt "i64"

let rec pp_ty fmt = function
  | TInt ty_int -> pp_ty_int fmt ty_int
  | TDouble -> fprintf fmt "double"
  | TPtr -> fprintf fmt "ptr"
  | TVoid -> fprintf fmt "void"
  | TArray (ty, size) ->
      if size = 0 then fprintf fmt "[%a]" pp_ty ty
      else fprintf fmt "[%d x %a]" size pp_ty ty

let pp_atom fmt atom =
  match (atom : atom) with
  | Int (ty_int, n) -> fprintf fmt "%a %Ld" pp_ty_int ty_int n
  | Double f -> fprintf fmt "%f" f
  | Var s -> fprintf fmt "%%%s" s

let pp_prim1 fmt (prim1 : prim1) =
  fprintf fmt "%s"
    (match prim1 with
    | Not -> "!"
    | Neg -> "-"
    | FNeg -> "-."
    | Float_of_int -> "float_of_int"
    | Int_of_float -> "int_of_float")

let pp_prim2 fmt (prim2 : prim2) =
  fprintf fmt "%s"
    (match prim2 with
    | Eq -> "="
    | Lt -> "<"
    | Le -> "<="
    | Gt -> ">"
    | Ge -> ">="
    | And -> "&"
    | Or -> "|"
    | Plus -> "+"
    | Minus -> "-"
    | Times -> "*"
    | Div -> "/"
    | Mod -> "mod"
    | Shl -> "shl"
    | Lshr -> "lshr"
    | Ashr -> "ashr"
    | FPlus -> "+."
    | FMinus -> "-."
    | FTimes -> "*."
    | FDiv -> "/."
    | FLt -> "<."
    | FLe -> "<=."
    | FEq -> "=."
    | FGe -> ">=."
    | FGt -> ">.")

let pp_instr name pp_args fmt args =
  fprintf fmt "@[%s %a@]" name (pp_list ~sep:",@;<1 2>" pp_args) args

let rec pp_expr fmt = function
  | Int (ty_int, n) -> pp_atom fmt (Int (ty_int, n))
  | Double f -> pp_atom fmt (Double f)
  | Var s -> pp_atom fmt (Var s)
  | Let (id, e1, e2) ->
      fprintf fmt "@[<v>@[<b>let@;<1 2>%s =@;<1 2>%a@;<1 0>in@]@;<0 0>%a@]" id
        pp_expr e1 pp_expr e2
  | Prim1 (op1, e) -> fprintf fmt "%a %a" pp_prim1 op1 pp_expr e
  | Prim2 (op2, e1, e2) ->
      fprintf fmt "@[(%a %a@;<1 2>%a)@]" pp_expr e1 pp_prim2 op2 pp_expr e2
  | If (e1, e2, e3) ->
      fprintf fmt "@[<hv>(if %a then@;<1 2>%a@;<1 1>else@;<1 2>%a)@]" pp_expr e1
        pp_expr e2 pp_expr e3
  | Switch (e, cases, default) ->
      fprintf fmt "switch %a with %a | default -> %a" pp_expr e
        (pp_list ~sep:"@ " (fun fmt (i, e) ->
             fprintf fmt "| %Ld -> %a" i pp_expr e))
        cases pp_expr default
  | For (id, e1, e2, e3) ->
      fprintf fmt "@[<v>@[for %s = %a to %a do@]@;<1 2>%a@;<1 0>done@]" id
        pp_expr e1 pp_expr e2 pp_expr e3
  | Apply (e_func, args, _) ->
      fprintf fmt "@[%a(@;<0 2>%a)@]" pp_expr e_func
        (pp_list ~sep:", " pp_expr)
        (List.map fst args)
  | ApplyExt (fname, _, _, exprs) ->
      fprintf fmt "%s(%a)" fname (pp_list ~sep:", " pp_expr) exprs
  | Load (ty, e) -> fprintf fmt "@[load@;<1 2>%a,@;<1 2>%a@]" pp_ty ty pp_expr e
  | Store (e1, e2) -> pp_instr "store" pp_expr fmt [ e1; e2 ]
  | Gep (ty, e, idx) ->
      fprintf fmt "@[gep@;<1 2>%a,@;<1 2>%a,@;<1 2>%d@]" pp_ty ty pp_expr e idx
  | Malloc ty -> pp_instr "malloc" pp_ty fmt [ ty ]
  | Alloca ty -> pp_instr "alloca" pp_ty fmt [ ty ]
  | Seq exprs ->
      fprintf fmt "@[<v>( %a)@]"
        (pp_list ~sep:";@;<0 2>" (fun fmt e -> fprintf fmt "@[%a@]" pp_expr e))
        exprs
  | Func fname -> fprintf fmt "func %s" fname
  | Extractelement (e_vec, e_idx) ->
      pp_instr "extractelement" pp_expr fmt [ e_vec; e_idx ]
  | Insertelement (e_vec, e_val, e_idx) ->
      pp_instr "insertelement" pp_expr fmt [ e_vec; e_val; e_idx ]
  | Inttoptr (e, ty) -> fprintf fmt "@[inttoptr %a to %a@]" pp_expr e pp_ty ty
  | Ptrtoint (e, ty) -> fprintf fmt "@[ptrtoint %a to %a@]" pp_expr e pp_ty ty
  | Intcast (e, ty) -> fprintf fmt "@[intcast %a to %a@]" pp_expr e pp_ty ty
  | Bitcast (e, ty) -> fprintf fmt "@[bitcast %a to %a@]" pp_expr e pp_ty ty

and pp_fundef fmt fundef =
  let pp_arg fmt (arg : string * ty) =
    fprintf fmt "(%s : %a)" (fst arg) pp_ty (snd arg)
  in
  fprintf fmt "@[<v>def %s %a : %a {@;<1 2>%a@;<0 0>}@]" fundef.name
    (pp_list ~sep:" " pp_arg) fundef.args pp_ty fundef.return pp_expr
    fundef.body

and pp_compunit fmt fundefs =
  fprintf fmt "@[<v>%a@]" (pp_list ~sep:"@;<0 0>@;<0 0>" pp_fundef) fundefs
