open Syntax
open Closure
open Common

(*

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`        | Values is the third LSB       |

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

Booleans are checked using direct bit patterns:
- `false` = `0x2` (`...0010`)
- `true`  = `0x6` (`...0110`)

### Boxed values

Boxed values are pointers to heap blocks. Normal pointers always have last two bits `00`

#### Mono backend:
| Value   | Block                             |
|---------|----------------------------------:|
| Tuple   | [ ev ; len ; vals ... ]           |
| Closure | [ ev ; fptr ; captured vals ... ] |
| Vec     | [ ev ; vals ... ]                 |
| Variant | [ ev ; id ; vals ... ]            |
| Float   | [ ev ; val ]                      |

#### Proxy backend:
| Value   | Block                             |
|---------|----------------------------------:|
| Tuple   | [ ev ; len ; vals ptr ]           |
| Closure | [ ev ; fptr ; captured vals ptr ] |
| Vec     | [ ev ; len ; vals ptr ]           |
| Variant | [ ev ; id ; vals ... ]            |
| Float   | [ ev ; val ]                      |

### Type discrimination

Value types can be determined as follows:
1. If LSB = 1: Integer
2. If LSB = 0:
   - If bit 1 = 1: Boolean (last two bits `10`)
   - Else: Boxed pointer (last two bits `00`)

*)
type t = Type.t

module Ev = Evidence.TypeEvidence

(* Backend selection flag *)
let proxy_values = ref false
let proxy_closures = ref false
let to_ptr t = Cmm.Syntax.Inttoptr (t, TPtr)
let bind x f b = Cmm.Syntax.Let (x, f, b)
let bind_list = List.fold_right2 bind

let load t idx =
  Cmm.Syntax.Load (TInt Int64, Gep (TInt Int64, Inttoptr (t, TPtr), idx))

let i64_t = Cmm.Syntax.TInt Int64
let unexpected_variant_case_error = 15L
let const_false = 0b010L
let const_true = 0b110L

(** Offsets *)

(** number of 8 byte words used by evidence *)
let evidence_word_num = 1

(* values *)

(** Mono: [ ev | len | vals ... ] Proxy: [ ev | len | vals_ptr ] *)
let tuple_offset = evidence_word_num + 1

(** Mono: [ ev | fptr | captured_vals ... ] Proxy:
    [ ev | fptr | captured_val_ptr ] *)
let closure_offset = evidence_word_num + 1

(** [ ev | case_id | vals ... ] *)
let variant_offset = evidence_word_num + 1

(* evidence *)

(** [ id | len | evs ... ] *)
let arrow_ev_offset = 2

(** [ id | len | evs ... ] *)
let tuple_ev_offset = 2

(** [ id | ev ] *)
let vec_ev_offset = 1

(* variant ids must be larger or equal than this number *)
let variant_id_base = 7
let variant_id ty = Type.variant_id ty + variant_id_base

let id_of_ty ty =
  match ty with
  | Type.Any -> 0
  | Type.Int -> 1
  | Type.Bool -> 2
  | Type.Vec _ -> 3
  | Type.Arrow _ -> 4
  | Type.Product _ -> 5
  | Type.Float -> 6
  | Type.Variant _ | Type.Rec _ -> variant_id ty
  | Type.Var _ ->
      failwith "this case should be avoided by the normalization phase"

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 initial_ev (const : Const.t) =
  match const with
  | Const.Int _ -> Type.Int
  | Const.Bool _ -> Type.Bool
  | Const.Float _ -> Type.Float

let heap_alloc_block num_slots =
  let size = num_slots * 8 in
  Cmm.Syntax.ApplyExt
    ("GC_malloc", [ TInt Int64 ], TPtr, [ Int (Int64, Int64.of_int size) ])

let heap_alloc ty =
  let size =
    match ty with
    | Cmm.Syntax.TArray (TInt Int64, n) -> n
    | _ -> failwith "not implemented yet"
  in
  heap_alloc_block size

let stack_alloc ty = Cmm.Syntax.Alloca ty

let allocate_block ?(alloc = heap_alloc) (terms : Cmm.Syntax.expr list)
    (k : Cmm.Syntax.expr -> Cmm.Syntax.expr) =
  let name = Gensym.fresh "block_ptr" in
  let block_name = Gensym.fresh "block" in
  let block_size = List.length terms in
  let block_type = Cmm.Syntax.TArray (TInt Int64, block_size) in
  let store_terms =
    List.mapi
      (fun i term -> Cmm.Syntax.Store (term, Gep (TInt Int64, Var name, i)))
      terms
  in
  Cmm.Syntax.Let
    ( name,
      alloc block_type,
      Seq
        (store_terms
        @ [
            Let (block_name, Ptrtoint (Var name, TInt Int64), k (Var block_name));
          ]) )

(* boxed values *)

let unbox_evidence t =
  Cmm.Syntax.Load (TDouble, Gep (TInt Int64, Inttoptr (t, TPtr), 0))

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

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

let rec lower_ev ?(alloc = heap_alloc) (ev : Type.t)
    (k : Cmm.Syntax.expr -> Cmm.Syntax.expr) =
  let ev_id = Cmm.Syntax.Int (Int64, Int64.of_int (id_of_ty ev)) in
  match ev with
  | Type.Any | Type.Int | Type.Bool | Type.Variant _ | Type.Rec _ | Type.Float
    ->
      k ev_id
  | Type.Vec ty ->
      lower_ev ~alloc ty (fun ev -> allocate_block ~alloc [ ev_id; ev ] k)
  | Type.Arrow (ty_args, ty_ret) ->
      let tys = ty_ret :: ty_args in
      let len = Cmm.Syntax.Int (Int64, Int64.of_int (List.length tys)) in
      let rec go acc tys =
        match tys with
        | [] -> allocate_block ~alloc (ev_id :: len :: List.rev acc) k
        | ty :: tys -> lower_ev ~alloc ty (fun ev -> go (ev :: acc) tys)
      in
      go [] tys
  | Type.Product tys ->
      let len = Cmm.Syntax.Int (Int64, Int64.of_int (List.length tys)) in
      let rec go acc tys =
        match tys with
        | [] -> allocate_block ~alloc (ev_id :: len :: List.rev acc) k
        | ty :: tys -> lower_ev ~alloc ty (fun ev -> go (ev :: acc) tys)
      in
      go [] tys
  | Type.Var _ ->
      failwith "this case should be avoided by the normalization phase"

(* Memory representation functions *)

let const (c : Const.t) : Cmm.Syntax.expr =
  match c with
  | Const.Int n -> Cmm.Syntax.Int (Int64, Int64.(add (mul (of_int n) 2L) 1L))
  | Const.Bool b -> Cmm.Syntax.Int (Int64, if b then const_true else const_false)
  | Const.Float f -> box_float (Cmm.Syntax.Double f)

let atom (atom : Type.t atom) : Cmm.Syntax.expr =
  match atom.atom_desc with Const c -> const c | Var x -> Cmm.Syntax.Var x

let allocate_closure fname captured_vars ev =
  let num_captured_vars =
    Cmm.Syntax.Int (Int64, Int64.of_int (List.length captured_vars))
  in
  lower_ev ev (fun ev ->
      Cmm.Syntax.ApplyExt
        ( "allocate_closure",
          [ TPtr; TInt Int64; TInt Int64 ],
          TInt Int64,
          [ Func fname; ev; num_captured_vars ] ))

let fill_closure closure_ptr _fname captured_vars =
  if List.length captured_vars = 0 then Cmm.Syntax.Int (Int64, 0L)
  else
    let closure_size =
      if !proxy_closures then 3 else closure_offset + List.length captured_vars
    in
    let closure =
      Cmm.Syntax.Gep
        (TArray (TInt Int64, closure_size), Inttoptr (closure_ptr, TPtr), 0)
    in
    if !proxy_closures then
      Cmm.Syntax.Let
        ( "free_vars_ptr",
          Load (TPtr, Gep (TInt Int64, closure, 2)),
          Seq
            (captured_vars
            |> List.mapi (fun i var ->
                   Cmm.Syntax.Store
                     ( Cmm.Syntax.Var var,
                       Gep (TInt Int64, Var "free_vars_ptr", i) ))) )
    else
      Cmm.Syntax.Seq
        (captured_vars
        |> List.mapi (fun i var ->
               let idx = closure_offset + i in
               let element = Cmm.Syntax.Gep (TInt Int64, closure, idx) in
               Cmm.Syntax.Store (Cmm.Syntax.Var var, element)))

let 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 ] ))

let bind_free_vars free_vars body =
  if !proxy_closures then
    bind "env_ptr" (load (Var "self") closure_offset)
    @@ bind_list free_vars
         (List.mapi (fun idx _ -> load (Var "env_ptr") idx) free_vars)
    @@ body
  else
    bind_list free_vars
      (List.mapi
         (fun idx _ -> load (Var "self") (closure_offset + idx))
         free_vars)
      body

let tuple (atoms : Type.t atom list) : Cmm.Syntax.expr =
  let open Cmm.Syntax in
  let len = List.length atoms in
  (* TODO: Refactor both implementations to increase code reuse *)
  if !proxy_values then
    let allocate_vals =
      let ty = TArray (TInt Int64, len) in
      heap_alloc ty
    in
    let fill_vals =
      let vals = Gep (TArray (TInt Int64, len), Var "vals", 0) in
      atoms
      |> List.mapi (fun i a ->
             let element = Gep (TInt Int64, vals, i) in
             Store (atom a, element))
    in
    let length = Int (Int64, Int64.of_int len) in
    Let
      ( "vals",
        allocate_vals,
        Seq
          (fill_vals
          @ [
              Let
                ( "ev",
                  ApplyExt
                    ( "make_tuple_ev",
                      [ TInt Int64; TPtr ],
                      TInt Int64,
                      [ length; Var "vals" ] ),
                  Let
                    ( "tuple",
                      ApplyExt
                        ( "make_tuple",
                          [ TInt Int64; TInt Int64; TPtr ],
                          TInt Int64,
                          [ Var "ev"; length; Var "vals" ] ),
                      Ptrtoint (Var "tuple", TInt Int64) ) );
            ]) )
  else
    let name = "tuple_ptr" in
    let size = tuple_offset + len in
    let allocate_tuple =
      let ty = TArray (TInt Int64, size) in
      heap_alloc ty
    in
    let fill_tuple =
      let tuple = Gep (TArray (TInt Int64, size), Var name, 0) in
      [
        (let element = Gep (TInt Int64, tuple, 1) in
         Store (Int (Int64, Int64.of_int len), element));
      ]
      @ (atoms
        |> List.mapi (fun i a ->
               let element = Gep (TInt Int64, tuple, tuple_offset + i) in
               Store (atom a, element)))
    in
    Let
      ( name,
        allocate_tuple,
        Seq
          (fill_tuple
          @ [
              ApplyExt
                ( "fill_tuple_ev",
                  [ TInt Int64 ],
                  TInt Int64,
                  [ Ptrtoint (Var name, TInt Int64) ] );
            ]
          @ [ Ptrtoint (Var name, TInt Int64) ]) )

let vect (a1 : Type.t atom) (a2 : Type.t atom) : Cmm.Syntax.expr =
  Cmm.Syntax.ApplyExt
    ("make_vector", [ TInt Int64; TInt Int64 ], TInt Int64, [ atom a1; atom a2 ])

let variant (case_name : string) (atoms : Type.t atom list) (ty : Type.t) :
    Cmm.Syntax.expr =
  let id = 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", variant_offset + i)
           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,
      Let
        ( "variant",
          Gep (TArray (TInt Int64, 0), Var "variant_ptr", 0),
          Seq (fill_variant @ [ Ptrtoint (Var "variant_ptr", TInt Int64) ]) ) )

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

let get (ptr : Cmm.Syntax.expr) (idx : int) : Cmm.Syntax.expr =
  Cmm.Syntax.ApplyExt
    ( "tuple_get",
      [ TInt Int64; TInt Int64 ],
      TInt Int64,
      [ ptr; Int (Int64, Int64.of_int idx) ] )

let ascribe (value : Cmm.Syntax.expr) (ty : Type.t) : Cmm.Syntax.expr =
  let alloc = stack_alloc in
  let is_any_list = List.for_all (( = ) Type.Any) in
  match ty with
  | Any -> value
  | Int ->
      Cmm.Syntax.ApplyExt ("ascribe_int", [ TInt Int64 ], TInt Int64, [ value ])
  | Bool ->
      Cmm.Syntax.ApplyExt ("ascribe_bool", [ TInt Int64 ], TInt Int64, [ value ])
  | Arrow (l, Any) when is_any_list l ->
      let length = 1 + List.length l in
      Cmm.Syntax.ApplyExt
        ( "ascribe_closure_dyn",
          [ TInt Int64; TInt Int64 ],
          TInt Int64,
          [ value; Int (Int64, Int64.of_int length) ] )
  | Arrow _ ->
      lower_ev ~alloc ty (fun ev ->
          Cmm.Syntax.ApplyExt
            ( "ascribe_closure",
              [ TInt Int64; TInt Int64 ],
              TInt Int64,
              [ value; ev ] ))
  | Vec Any ->
      Cmm.Syntax.ApplyExt
        ("ascribe_vec_dyn", [ TInt Int64 ], TInt Int64, [ value ])
  | Vec _ ->
      lower_ev ~alloc ty (fun ev ->
          Cmm.Syntax.ApplyExt
            ( "ascribe_vec",
              [ TInt Int64; TInt Int64 ],
              TInt Int64,
              [ value; ev ] ))
  | Product l when is_any_list l ->
      let length = List.length l in
      Cmm.Syntax.ApplyExt
        ( "ascribe_tuple_dyn",
          [ TInt Int64; TInt Int64 ],
          TInt Int64,
          [ value; Int (Int64, Int64.of_int length) ] )
  | Product _ ->
      lower_ev ~alloc ty (fun ev ->
          Cmm.Syntax.ApplyExt
            ( "ascribe_tuple",
              [ TInt Int64; TInt Int64 ],
              TInt Int64,
              [ value; ev ] ))
  | Variant _ | Rec (_, _) | Float ->
      lower_ev ~alloc ty (fun ev ->
          Cmm.Syntax.ApplyExt
            ( "ascribe_imm_evidence",
              [ TInt Int64; TInt Int64 ],
              TInt Int64,
              [ value; ev ] ))
  | Var _ -> failwith "this case should be avoided by the normalization phase"

let ite (c_val : Cmm.Syntax.expr) (true_e : Cmm.Syntax.expr)
    (false_e : Cmm.Syntax.expr) : Cmm.Syntax.expr =
  Cmm.Syntax.If (Intcast (untag_bool c_val, TInt Int1), true_e, false_e)

let for_ (var : string) (start : Cmm.Syntax.expr) (stop : Cmm.Syntax.expr)
    (body : Cmm.Syntax.expr) : Cmm.Syntax.expr =
  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 (a1 : Cmm.Syntax.expr) (a2 : Cmm.Syntax.expr) : Cmm.Syntax.expr =
  Cmm.Syntax.ApplyExt
    ("vector_get", [ TInt Int64; TInt Int64 ], TInt Int64, [ a1; a2 ])

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

let vect_length (a1 : Cmm.Syntax.expr) : Cmm.Syntax.expr =
  Cmm.Syntax.ApplyExt ("vector_length", [ TInt Int64 ], TInt Int64, [ a1 ])

let bind_args variant args body =
  let offset = variant_offset in
  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, offset + idx),
             Let (name, Load (TInt Int64, Var slot_name), acc) ))
       body

let match_ (variant_ptr : Cmm.Syntax.expr)
    (cases : (string * string list * Cmm.Syntax.expr) list) (_ty : Type.t) :
    Cmm.Syntax.expr =
  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_ptr",
      Inttoptr (variant_ptr, TPtr),
      Let
        ( "variant",
          Gep (TArray (TInt Int64, 0), Var "variant_ptr", 0),
          Let
            ( "case_id",
              Load (TInt Int64, Gep (TInt Int64, Var "variant", 1)),
              Switch (Var "case_id", cases (Var "variant"), default) ) ) )

let rec codegen_program (program : Type.t program) : Cmm.Syntax.compunit =
  let funcs = List.map codegen_func program.funcs in
  let func =
    Cmm.Syntax.
      {
        name = "main";
        args = [];
        return = TInt Int32;
        body = Seq [ codegen_expr program.expr; Int (Int32, 0L) ];
      }
  in
  funcs @ [ func ]

and codegen_expr (expr : Type.t expr) : Cmm.Syntax.expr =
  match expr.expr_desc with
  | Ret cexpr -> codegen_cexpr cexpr
  | Let (binding, expr) ->
      Cmm.Syntax.Let
        (binding.name, codegen_cexpr binding.cexpr, codegen_expr expr)
  | Letrec (bindings, expr) ->
      (* fail if expressions other than lambdas are bound with letrec *)
      let closure_bindings =
        bindings
        |> List.map (fun binding ->
               match binding.cexpr.cexpr_desc with
               | Closure (fname, captured_vars) ->
                   ( binding.name,
                     (fname, captured_vars),
                     binding.cexpr.cexpr_meta )
               | _ -> failwith "expression other than lambda bound with let rec")
      in
      (* allocate bound closures *)
      let rec allocate_closures l ret =
        match l with
        | (id, (fname, captured_vars), ev) :: tl ->
            Cmm.Syntax.Let
              ( id,
                allocate_closure fname captured_vars ev,
                allocate_closures tl ret )
        | [] -> ret
      in
      (* fill allocated closures *)
      let fill_closures =
        closure_bindings
        |> List.map (fun (id, (fname, captured_vars), _) ->
               fill_closure Cmm.Syntax.(Var id) fname captured_vars)
      in
      allocate_closures closure_bindings
        (Cmm.Syntax.Seq (fill_closures @ [ codegen_expr expr ]))
  | Seq exprs -> Cmm.Syntax.Seq (List.map codegen_expr exprs)

and codegen_cexpr (cexpr : Type.t cexpr) : Cmm.Syntax.expr =
  match cexpr.cexpr_desc with
  | Atom a -> codegen_atom a
  | Call (closure, args, well_known) -> call ~well_known closure args
  | Closure (fname, free_vars) -> closure fname free_vars cexpr.cexpr_meta
  | Prim1 (op1, c1) -> prim1 op1 (codegen_cexpr c1)
  | Prim2 (op2, a1, a2) -> prim2 op2 (codegen_atom a1) (codegen_atom a2)
  | If (a, e1, e2) -> ite (codegen_atom a) (codegen_expr e1) (codegen_expr e2)
  | Match (a, patterns, ty) ->
      let a = codegen_atom a in
      let patterns =
        List.map
          (fun Closure.{ casename; args; body } ->
            (casename, args, codegen_expr body))
          patterns
      in
      match_ a patterns ty
  | Tuple atoms -> tuple atoms
  | Get (c, i) -> get (codegen_cexpr c) i
  | Vect (a1, a2) -> vect a1 a2
  | VectGet (a1, a2) -> vect_get (codegen_atom a1) (codegen_atom a2)
  | VectSet (vector_a, index_a, value_a) -> (
      let vector = codegen_atom vector_a in
      let index = codegen_atom index_a in
      let value = codegen_atom value_a in
      let vec_ev = vector_a.atom_meta in
      let value_ev = value_a.atom_meta in
      match Evidence.TypeMatching.vec vec_ev with
      | Some ev
        when Ev.is_fully_precise value_ev
             && Ev.is_fully_precise ev
             && value_ev = ev ->
          vect_set vector index value
      | _ ->
          bind "ev" (load vector 0)
          @@ bind "ev_inner" (load (Var "ev") vec_ev_offset)
          @@ bind "value"
               (Cmm.Syntax.ApplyExt
                  ("ascribe", [ i64_t; i64_t ], i64_t, [ value; Var "ev_inner" ]))
          @@ vect_set vector index (Var "value"))
  | VectLength c -> vect_length (codegen_cexpr c)
  | Ascribe (c, ty) -> ascribe (codegen_cexpr c) ty
  | AscribeDyn (c, a, inv) -> ascribe_dyn (codegen_cexpr c) (codegen_atom a) inv
  | For (var, start, stop, body) ->
      for_ var (codegen_atom start) (codegen_atom stop) (codegen_expr body)
  | Variant (case_name, atoms, ty) -> variant case_name atoms ty

and codegen_atom (atom : Type.t atom) : Cmm.Syntax.expr =
  match atom.atom_desc with Const c -> const c | Var x -> Cmm.Syntax.Var x

and codegen_func (func : 'a func) : Cmm.Syntax.fundef =
  match func with
  | Func f ->
      let body = bind_free_vars f.free_vars (codegen_expr f.body) in
      Cmm.Syntax.
        {
          name = f.name;
          args = ("self", TPtr) :: List.map (fun x -> (x, TInt Int64)) f.args;
          return = TInt Int64;
          body;
        }

and call ?(well_known = None) closure args : Cmm.Syntax.expr =
  let open Cmm.Syntax in
  let arg_names = List.mapi (fun i _ -> Format.sprintf "arg_%d" i) args in
  let args' = List.map (fun s -> (Var s, i64_t)) arg_names in
  let func =
    match well_known with Some fname -> Func fname | None -> Var "f_ptr"
  in
  let lowered_args = List.map codegen_atom args in
  let closure = codegen_atom closure in
  bind "closure_ptr" (to_ptr closure)
  @@ bind "evidence_ptr" (to_ptr (load (Var "closure_ptr") 0))
  @@ bind "f_ptr" (to_ptr (load (Var "closure_ptr") 1))
  @@ bind_list arg_names lowered_args
  @@ Apply (func, (Var "closure_ptr", TPtr) :: args', i64_t)

and ascribe_dyn value arg inv =
  match inv with
  | Inversion.IDom i ->
      bind "closure_ptr" (to_ptr arg)
      @@ bind "evidence_ptr" (to_ptr (load (Var "closure_ptr") 0))
      @@ bind "ev" (load (Var "evidence_ptr") (arrow_ev_offset + 1 + i))
      @@ ApplyExt ("ascribe", [ i64_t; i64_t ], i64_t, [ value; Var "ev" ])
  | Inversion.IProd i ->
      bind "tuple_ptr" (to_ptr arg)
      @@ bind "evidence_ptr" (to_ptr (load (Var "tuple_ptr") 0))
      @@ ApplyExt
           ( "ascribe",
             [ i64_t; i64_t ],
             i64_t,
             [ value; load (Var "evidence_ptr") (tuple_ev_offset + i) ] )
  | Inversion.ICod ->
      bind "closure_ptr" (to_ptr arg)
      @@ bind "evidence_ptr" (to_ptr (load (Var "closure_ptr") 0))
      @@ ApplyExt
           ( "ascribe",
             [ i64_t; i64_t ],
             i64_t,
             [ value; load (Var "evidence_ptr") arrow_ev_offset ] )
  | Inversion.IVec _ ->
      bind "vec_ptr" (to_ptr arg)
      @@ bind "evidence_ptr" (to_ptr (load (Var "vec_ptr") 0))
      @@ Cmm.Syntax.ApplyExt
           ( "ascribe",
             [ i64_t; i64_t ],
             i64_t,
             [ value; load (Var "evidence_ptr") vec_ev_offset ] )
