open Syntax
open Closure
open Lower

let unexpected_variant_case_error = 15L

(*

Values are either immediate values or boxed. The least significant bit (LSB)
indicates whether a value is immediate or boxed.

| Type      | Representation |
|-----------|---------------:|
| Immediate |  `B...B1`   |
| Boxed     |  `B...B0`   |

### Immediate values

Immediate values include integers and booleans.

| Value   | Representation      | Description                   |
|---------|--------------------:|-------------------------------|
| int     | `BBB...BBB1`        | Shifted left by 1, LSB=1      |
| bool    | `000...0B10`        | Value is the fourth LSB       |

Integer arithmetic operations must account for this representation:
- `x + y` → `x + y - 2`
- `x * y` → `((x / 2) * (y / 2)) * 2 + 1`

Booleans are exactly the following values:
- `false` = `14` (`0b01110`)
- `true`  = `30` (`0b11110`)

### Boxed values

Boxed values are stored in the heap, and their first 64 bits are
a tag indicating the type of value.

| Value   | Tag    |
|---------|-------:|
| Tuple   | `0010` |
| Closure | `0100` |
| Float   | `0110` | 
| Vector  | `1000` |
| Variant | `1010` |

### Type discrimination

Value types can be determined as follows:
1. If LSB = 1: Integer
2. If LSB = 0:
   - If suffix is 0b1110: Boolean
   - Else: Boxed pointer
*)

(* offsets *)
let closure_vars_offset = 2
let tuple_vars_offset = 2
let vec_vars_offset = 2
let variant_vars_offset = 3
let variant_caseid_offset = 2

(* bools *)

let const_false = 0b010L
let const_true = 0b110L

(* tags for boxed values *)
let tuple_tag = 3L
let closure_tag = 4L
let float_tag = 5L
let vector_tag = 6L
let variant_tag = 7L

let tag_int t =
  Cmm.Syntax.Prim2 (Or, Prim2 (Shl, t, Int (Int64, 1L)), Int (Int64, 1L))

let untag_int t = Cmm.Syntax.Prim2 (Ashr, t, Int (Int64, 1L))

let tag_bool t =
  Cmm.Syntax.Prim2
    (Or, Prim2 (Times, t, Int (Int64, 4L)), Int (Int64, const_false))

let untag_bool t =
  Cmm.Syntax.Prim2
    (Div, Prim2 (Minus, t, Int (Int64, const_false)), Int (Int64, 4L))

let unbox_float f =
  Cmm.Syntax.Load (TDouble, Gep (TDouble, Inttoptr (f, TPtr), 1))

let box_float f =
  Cmm.Syntax.ApplyExt ("make_float", [ TDouble ], TInt Int64, [ f ])

module TaggedMemrep :
  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 -> tag_int @@ Cmm.Syntax.Int (Int64, Int64.of_int n)
    | Const.Bool b ->
        Cmm.Syntax.Int (Int64, if b then const_true else const_false)
    | Const.Float f -> box_float @@ Cmm.Syntax.Double f

  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 closure_size = closure_vars_offset + List.length captured_vars in
    let block_size = closure_size * 8 in
    Cmm.Syntax.Ptrtoint
      ( ApplyExt
          ( "GC_malloc",
            [ TInt Int64 ],
            TPtr,
            [ Int (Int64, Int64.of_int block_size) ] ),
        TInt Int64 )

  and fill_closure closure_ptr fname captured_vars =
    let f = Cmm.Syntax.Func fname in
    let closure =
      Cmm.Syntax.Gep (TInt Int64, Inttoptr (closure_ptr, TPtr), 0)
    in
    let first_slot = Cmm.Syntax.Gep (TInt Int64, closure, 0) in
    let second_slot = Cmm.Syntax.Gep (TInt Int64, closure, 1) in
    Cmm.Syntax.Seq
      (Store (Int (Int64, closure_tag), first_slot)
      :: Store (f, second_slot)
      :: (captured_vars
         |> List.mapi (fun i var ->
                let idx = i + closure_vars_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 =
    Cmm.Syntax.(
      Let
        ( "closure_ptr",
          allocate_closure fname captured_vars ev,
          Seq [ fill_closure (Var "closure_ptr") fname captured_vars ] ))

  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,
                   Inttoptr (Var "self", TPtr),
                   idx + closure_vars_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 len = List.length atoms in
    let size = len + tuple_vars_offset 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 =
      let tuple = Cmm.Syntax.Gep (TInt Int64, Var name, 0) in
      let first_slot = Cmm.Syntax.Gep (TInt Int64, tuple, 0) in
      let second_slot = Cmm.Syntax.Gep (TInt Int64, tuple, 1) in
      [
        Cmm.Syntax.Store (Int (Int64, tuple_tag), first_slot);
        Cmm.Syntax.Store (Int (Int64, Int64.of_int len), second_slot);
      ]
      @ (atoms
        |> List.mapi (fun i a ->
               let element =
                 Cmm.Syntax.Gep (TInt Int64, tuple, i + tuple_vars_offset)
               in
               Cmm.Syntax.Store (atom a, element)))
    in
    Cmm.Syntax.Let
      ( name,
        allocate_tuple,
        Seq (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_ptr", i + variant_vars_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_ptr",
        allocate_variant,
        Seq (fill_variant @ [ Ptrtoint (Var "variant_ptr", TInt Int64) ]) )
end

module TaggedElim : 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 -> tag_bool (Prim1 (Not, untag_bool 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, [ value ])
        | 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 -> box_float @@ Cmm.Syntax.Prim1 (FNeg, unbox_float value)
    | Prim.Float_of_int -> box_float @@ Prim1 (Float_of_int, untag_int value)
    | Prim.Int_of_float ->
        tag_int (Cmm.Syntax.Prim1 (Int_of_float, unbox_float value))
    | Prim.FSqrt | Prim.FSin | Prim.FCos | Prim.FExp | Prim.FLog | Prim.FRound
      ->
        let fname = Format.asprintf "%a" Prim.pp_op1 op1 in
        box_float
        @@ Cmm.Syntax.ApplyExt
             (fname, [ TDouble ], TDouble, [ unbox_float value ])

  let prim2 : Prim.op2 -> t -> t -> t =
   fun op2 v1 v2 ->
    let i64_of_i1 t =
      Cmm.Syntax.Prim2 (And, Intcast (t, TInt Int64), Int (Int64, 1L))
    in
    let lift_float2 op v1 v2 =
      box_float @@ Prim2 (op, unbox_float v1, unbox_float v2)
    in
    let lift_float_cmp op v1 v2 =
      tag_bool @@ i64_of_i1 @@ Prim2 (op, unbox_float v1, unbox_float v2)
    in
    match op2 with
    | Prim.Plus -> Prim2 (Plus, Prim2 (Plus, v1, v2), Int (Int64, -1L))
    | Prim.Minus -> Prim2 (Plus, Prim2 (Minus, v1, v2), Int (Int64, 1L))
    | Prim.Times -> tag_int @@ Prim2 (Times, untag_int v1, untag_int v2)
    | Prim.Div -> tag_int @@ Prim2 (Div, untag_int v1, untag_int v2)
    | Prim.Mod -> tag_int @@ Prim2 (Mod, untag_int v1, untag_int v2)
    | Prim.And -> tag_bool @@ Prim2 (And, untag_bool v1, untag_bool v2)
    | Prim.Or -> tag_bool @@ Prim2 (Or, untag_bool v1, untag_bool v2)
    | Prim.Eq -> tag_bool @@ i64_of_i1 @@ Prim2 (Eq, v1, v2)
    | Prim.Less -> tag_bool @@ i64_of_i1 @@ Prim2 (Lt, v1, v2)
    | Prim.LessEq -> tag_bool @@ i64_of_i1 @@ Prim2 (Le, v1, v2)
    | Prim.Greater -> tag_bool @@ i64_of_i1 @@ Prim2 (Gt, v1, v2)
    | Prim.GreaterEq -> tag_bool @@ i64_of_i1 @@ Prim2 (Ge, 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 -> lift_float_cmp FEq v1 v2
    | Prim.FLess -> lift_float_cmp FLt v1 v2
    | Prim.FLessEq -> lift_float_cmp FLe v1 v2
    | Prim.FGreater -> lift_float_cmp FGt v1 v2
    | Prim.FGreaterEq -> lift_float_cmp FGe v1 v2
    | Prim.FMax ->
        box_float
        @@ ApplyExt
             ( "fmax",
               [ TDouble; TDouble ],
               TDouble,
               [ unbox_float v1; unbox_float v2 ] )
    | Prim.FMin ->
        box_float
        @@ ApplyExt
             ( "fmin",
               [ TDouble; TDouble ],
               TDouble,
               [ unbox_float v1; unbox_float v2 ] )

  let get : t -> int -> t =
   fun tuple idx ->
    Cmm.Syntax.ApplyExt
      ( "tuple_get",
        [ TInt Int64; TInt Int64 ],
        TInt Int64,
        [ tuple; Int (Int64, Int64.of_int idx) ] )

  let ascribe : t -> Type.t -> t =
   fun t ty ->
    match ty with
    | Type.Int ->
        Seq [ ApplyExt ("assert_int", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Bool ->
        Seq [ ApplyExt ("assert_bool", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Float ->
        Seq [ ApplyExt ("assert_float", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Any -> t
    | Type.Arrow (_, _) ->
        Seq [ ApplyExt ("assert_closure", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Product _ ->
        Seq [ ApplyExt ("assert_tuple", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Vec _ ->
        Seq [ ApplyExt ("assert_vec", [ TInt Int64 ], TVoid, [ t ]); t ]
    | Type.Variant _ | Type.Rec (_, _) ->
        let variant_id = Type.variant_id ty in
        Seq
          [
            ApplyExt
              ( "assert_variant",
                [ TInt Int64; TInt Int64 ],
                TVoid,
                [ t; Int (Int64, Int64.of_int variant_id) ] );
            t;
          ]
    | Type.Var _ -> failwith "should be impossible"

  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" closure
    @@ bind "closure_ptr" (Inttoptr (Var "closure", TPtr))
    @@ bind "f_ptr"
         (to_ptr (Load (TInt Int64, Gep (TInt Int64, Var "closure_ptr", 1))))
    @@ Apply (func, (Var "closure_ptr", TPtr) :: args, return_type)

  let ite : t -> t -> t -> t =
   fun c_val true_e false_e ->
    Cmm.Syntax.If (Intcast (untag_bool c_val, TInt Int1), true_e, false_e)

  let for_ : string -> t -> t -> t -> t =
   fun var start stop body ->
    let var' = var ^ "_i" in
    Cmm.Syntax.For
      ( var',
        untag_int start,
        untag_int stop,
        Let (var, tag_int (Var var'), 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_vars_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 err_code = Cmm.Syntax.Int (Int64, unexpected_variant_case_error) in
    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 =
      let noop = Cmm.Syntax.Int (Int64, 0L) in
      Cmm.Syntax.Seq
        [ ApplyExt ("error", [ TInt Int64 ], TVoid, [ err_code ]); noop ]
    in
    Cmm.Syntax.Let
      ( "variant",
        Inttoptr (variant_ptr, TPtr),
        Let
          ( "case_id",
            Load
              ( TInt Int64,
                Gep (TInt Int64, Var "variant", variant_caseid_offset) ),
            Switch (Var "case_id", cases (Var "variant"), default) ) )
end

module Lower = MakeLower (TaggedMemrep) (TaggedElim)
