open Syntax
open Closure
open Lower

(** [ vals ... ] *)
let tuple_offset = 0

(** [ fptr | captured_vals ... ] *)
let closure_offset = 1

(** [ id | case_id | vals ... ] *)
let variant_offset = 2

let unexpected_variant_case_error = 15L

module UntaggedMemrep :
  MemRep with type t = Cmm.Syntax.expr and type ev = Type.t = struct
  type t = Cmm.Syntax.expr
  type ev = Type.t

  let rec const : Const.t -> t = function
    | Const.Int n -> Cmm.Syntax.Int (Int64, Int64.of_int n)
    | Const.Bool b ->
        Cmm.Syntax.Int (Int64, match b with false -> 0L | true -> 1L)
    | Const.Float f -> Cmm.Syntax.Bitcast (Double f, TInt Int64)

  and atom : ev atom -> t =
   fun atom ->
    match atom.atom_desc with Const c -> const c | Var x -> Cmm.Syntax.Var x

  and allocate_closure fname captured_vars _ =
    let num_captured_vars =
      Cmm.Syntax.Int (Int64, Int64.of_int (List.length captured_vars))
    in
    Cmm.Syntax.(
      ApplyExt
        ( "allocate_closure",
          [ TPtr; TInt Int64 ],
          TInt Int64,
          [ Func fname; num_captured_vars ] ))

  and fill_closure closure_ptr fname captured_vars =
    let f = Cmm.Syntax.Func fname in
    let closure = Cmm.Syntax.Inttoptr (closure_ptr, TPtr) in
    let first_slot = Cmm.Syntax.Gep (TInt Int64, closure, 0) in
    Cmm.Syntax.Seq
      (Store (f, first_slot)
      :: (captured_vars
         |> List.mapi (fun i var ->
                let idx = i + closure_offset in
                let element = Cmm.Syntax.Gep (TInt Int64, closure, idx) in
                Cmm.Syntax.Store (Cmm.Syntax.Var var, element))))

  and closure fname captured_vars ev =
    let name = Format.sprintf "closure_%s" fname in
    Cmm.Syntax.(
      Let
        ( name,
          allocate_closure fname captured_vars ev,
          Seq [ fill_closure (Var name) fname captured_vars; Var name ] ))

  and bind_free_vars free_vars body =
    free_vars
    |> List.mapi (fun idx name -> (idx, name))
    |> List.fold_left
         (fun acc (idx, name) ->
           Cmm.Syntax.Let
             ( "ptr",
               Gep (TInt Int64, Var "self", idx + closure_offset),
               Let (name, Load (TInt Int64, Var "ptr"), acc) ))
         body

  and tuple : 'a atom list -> Cmm.Syntax.expr =
   fun atoms ->
    let name = "tuple_ptr" in
    let size = List.length atoms in
    let allocate_tuple =
      let block_size = size * 8 in
      Cmm.Syntax.ApplyExt
        ( "GC_malloc",
          [ TInt Int64 ],
          TPtr,
          [ Int (Int64, Int64.of_int block_size) ] )
    in
    let fill_tuple =
      atoms
      |> List.mapi (fun i a ->
             let element =
               Cmm.Syntax.Gep (TInt Int64, Var "tuple_ptr", i + tuple_offset)
             in
             Cmm.Syntax.Store (atom a, element))
    in
    Cmm.Syntax.Let
      ( name,
        allocate_tuple,
        Seq (List.append fill_tuple [ Ptrtoint (Var name, TInt Int64) ]) )

  and vect : 'a atom -> 'a atom -> Cmm.Syntax.expr =
   fun a1 a2 ->
    Cmm.Syntax.ApplyExt
      ( "make_vector",
        [ TInt Int64; TInt Int64 ],
        TInt Int64,
        [ atom a1; atom a2 ] )

  and variant : string -> ev atom list -> Type.t -> Cmm.Syntax.expr =
   fun case_name atoms ty ->
    let id = Type.variant_id ty in
    let case_id = Type.case_id case_name ty in
    let num_args = List.length atoms in
    let fill_variant =
      atoms
      |> List.mapi (fun i a ->
             let element =
               Cmm.Syntax.Gep (TInt Int64, Var "variant", i + variant_offset)
             in
             Cmm.Syntax.Store (atom a, element))
    in
    let allocate_variant =
      Cmm.Syntax.Inttoptr
        ( ApplyExt
            ( "make_variant",
              [ TInt Int64; TInt Int64; TInt Int64 ],
              TInt Int64,
              [
                Int (Int64, Int64.of_int id);
                Int (Int64, Int64.of_int case_id);
                Int (Int64, Int64.of_int num_args);
              ] ),
          TPtr )
    in
    Cmm.Syntax.Let
      ( "variant",
        allocate_variant,
        Seq (fill_variant @ [ Ptrtoint (Var "variant", TInt Int64) ]) )
end

module UntaggedElim : Elim with type t = Cmm.Syntax.expr and type ev = unit =
struct
  type t = Cmm.Syntax.expr
  type ev = unit

  let prim1 : Prim.op1 -> t -> t =
   fun op1 value ->
    match op1 with
    | Prim.Neg -> Cmm.Syntax.(Prim1 (Neg, value))
    | Prim.Not -> Cmm.Syntax.(Prim1 (Not, value))
    | Prim.Print print -> (
        match print with
        | Prim.PrintInt ->
            Cmm.Syntax.ApplyExt
              ("print_int", [ TInt Int64 ], TInt Int64, [ value ])
        | Prim.PrintBool ->
            Cmm.Syntax.ApplyExt
              ("print_bool", [ TInt Int64 ], TInt Int64, [ value ])
        | Prim.PrintFloat ->
            Cmm.Syntax.ApplyExt
              ( "print_float",
                [ TInt Int64 ],
                TInt Int64,
                [ Bitcast (value, TInt Int64) ] )
        | Prim.PrintAny ->
            Cmm.Syntax.ApplyExt ("print", [ TInt Int64 ], TInt Int64, [ value ])
        | Prim.PrintChar c ->
            Cmm.Syntax.ApplyExt
              ( "print_char",
                [ TInt Int64 ],
                TInt Int64,
                [ Int (Int64, Int64.of_int (Char.code c)) ] ))
    | Prim.ReadInt -> Cmm.Syntax.ApplyExt ("read_int", [], TInt Int64, [])
    | Prim.ReadFloat -> Cmm.Syntax.ApplyExt ("read_float", [], TInt Int64, [])
    | Prim.ReadChar -> Cmm.Syntax.ApplyExt ("read_char", [], TInt Int64, [])
    | Prim.Time ->
        Cmm.Syntax.ApplyExt
          ("print_time", [ TPtr ], TInt Int64, [ Inttoptr (value, TPtr) ])
    | Prim.FNeg -> Bitcast (Prim1 (FNeg, Bitcast (value, TDouble)), TInt Int64)
    | Prim.Float_of_int -> Prim1 (Float_of_int, value)
    | Prim.Int_of_float ->
        Bitcast (Prim1 (Int_of_float, Bitcast (value, TDouble)), TInt Int64)
    | Prim.FSqrt | Prim.FSin | Prim.FCos | Prim.FExp | Prim.FLog | Prim.FRound
      ->
        let fname = Format.asprintf "%a" Prim.pp_op1 op1 in
        Bitcast
          ( ApplyExt (fname, [ TDouble ], TDouble, [ Bitcast (value, TDouble) ]),
            TInt Int64 )

  let prim2 : Prim.op2 -> t -> t -> t =
   fun op2 v1 v2 ->
    let lift_float2 op v1 v2 =
      Cmm.Syntax.Bitcast
        (Prim2 (op, Bitcast (v1, TDouble), Bitcast (v2, TDouble)), TInt Int64)
    in
    let i64_of_i1 v =
      Cmm.Syntax.Prim2 (And, Intcast (v, TInt Int64), Int (Int64, 1L))
    in
    match op2 with
    | Prim.Plus -> Prim2 (Plus, v1, v2)
    | Prim.Minus -> Prim2 (Minus, v1, v2)
    | Prim.Times -> Prim2 (Times, v1, v2)
    | Prim.And -> i64_of_i1 @@ Prim2 (And, v1, v2)
    | Prim.Or -> i64_of_i1 @@ Prim2 (Or, v1, v2)
    | Prim.Eq -> i64_of_i1 @@ Prim2 (Eq, v1, v2)
    | Prim.Less -> i64_of_i1 @@ Prim2 (Lt, v1, v2)
    | Prim.LessEq -> i64_of_i1 @@ Prim2 (Le, v1, v2)
    | Prim.Greater -> i64_of_i1 @@ Prim2 (Gt, v1, v2)
    | Prim.GreaterEq -> i64_of_i1 @@ Prim2 (Ge, v1, v2)
    | Prim.Div -> Prim2 (Div, v1, v2)
    | Prim.Mod -> Prim2 (Mod, v1, v2)
    | Prim.FPlus -> lift_float2 FPlus v1 v2
    | Prim.FMinus -> lift_float2 FMinus v1 v2
    | Prim.FTimes -> lift_float2 FTimes v1 v2
    | Prim.FDiv -> lift_float2 FDiv v1 v2
    | Prim.FEq ->
        i64_of_i1 @@ Prim2 (FEq, Bitcast (v1, TDouble), Bitcast (v2, TDouble))
    | Prim.FLess ->
        i64_of_i1 @@ Prim2 (FLt, Bitcast (v1, TDouble), Bitcast (v2, TDouble))
    | Prim.FLessEq ->
        i64_of_i1 @@ Prim2 (FLe, Bitcast (v1, TDouble), Bitcast (v2, TDouble))
    | Prim.FGreater ->
        i64_of_i1 @@ Prim2 (FGt, Bitcast (v1, TDouble), Bitcast (v2, TDouble))
    | Prim.FGreaterEq ->
        i64_of_i1 @@ Prim2 (FGe, Bitcast (v1, TDouble), Bitcast (v2, TDouble))
    | Prim.FMax ->
        Bitcast
          ( ApplyExt
              ( "fmax",
                [ TDouble; TDouble ],
                TDouble,
                [ Bitcast (v1, TDouble); Bitcast (v2, TDouble) ] ),
            TInt Int64 )
    | Prim.FMin ->
        Bitcast
          ( ApplyExt
              ( "fmin",
                [ TDouble; TDouble ],
                TDouble,
                [ Bitcast (v1, TDouble); Bitcast (v2, TDouble) ] ),
            TInt Int64 )

  let get : t -> int -> t =
   fun ptr idx ->
    Cmm.Syntax.(
      Load
        (TInt Int64, Gep (TInt Int64, Inttoptr (ptr, TPtr), tuple_offset + idx)))

  let ascribe : t -> Type.t -> t = fun value _ -> value

  let call : ?well_known:string option -> t -> t list -> t =
   fun ?(well_known = None) closure args ->
    let open Cmm.Syntax in
    let to_ptr t = Inttoptr (t, TPtr) in
    let bind x f b = Let (x, f, b) in
    let args = List.map (fun arg -> (arg, TInt Int64)) args in
    let return_type = TInt Int64 in
    let func =
      match well_known with Some fname -> Func fname | None -> Var "f_ptr"
    in
    bind "closure_ptr" (to_ptr closure)
    @@ bind "f_ptr"
         (to_ptr (Load (TInt Int64, Gep (TInt Int64, Var "closure_ptr", 0))))
    @@ Apply (func, (Var "closure_ptr", TPtr) :: args, return_type)

  let ite : t -> t -> t -> t =
   fun cond true_ false_ ->
    Cmm.Syntax.If (Intcast (cond, TInt Int1), true_, false_)

  let for_ : string -> t -> t -> t -> t =
   fun var start stop body -> Cmm.Syntax.For (var, start, stop, body)

  let vect_get : t -> t -> t =
   fun a1 a2 ->
    Cmm.Syntax.ApplyExt
      ("vector_get", [ TInt Int64; TInt Int64 ], TInt Int64, [ a1; a2 ])

  let vect_set : t -> t -> t -> t =
   fun a1 a2 a3 ->
    Cmm.Syntax.ApplyExt
      ( "vector_set",
        [ TInt Int64; TInt Int64; TInt Int64 ],
        TInt Int64,
        [ a1; a2; a3 ] )

  let vect_length : t -> t =
   fun a1 ->
    Cmm.Syntax.ApplyExt ("vector_length", [ TInt Int64 ], TInt Int64, [ a1 ])

  let bind_args variant args body =
    args
    |> List.mapi (fun idx name -> (idx, name))
    |> List.fold_left
         (fun acc (idx, name) ->
           let slot_name = Format.sprintf "slot_%s" name in
           Cmm.Syntax.Let
             ( slot_name,
               Gep (TInt Int64, variant, idx + variant_offset),
               Let (name, Load (TInt Int64, Var slot_name), acc) ))
         body

  let match_ : t -> (string * string list * t) list -> Type.t -> t =
   fun variant_ptr cases _ty ->
    let cases variant =
      cases
      |> List.mapi (fun i (_, args, body) ->
             let case_id = Int64.of_int i in
             let body = bind_args variant args body in
             (case_id, body))
    in
    let default = Cmm.Syntax.Int (Int64, 0L) in
    Cmm.Syntax.Let
      ( "variant",
        Inttoptr (variant_ptr, TPtr),
        Let
          ( "case_id",
            Load (TInt Int64, Gep (TInt Int64, Var "variant", 1)),
            Switch (Var "case_id", cases (Var "variant"), default) ) )
end

module Lower = MakeLower (UntaggedMemrep) (UntaggedElim)
