module SuccinctTree = struct
  type t = {
    topol : int64;
    labels : int64;
    topol_length : int;
    labels_length : int;
    width : int;
    alphabet_size : int;
  }

  let max_width = 64
  let lparen = 1L
  let rparen = 0L

  let width ~alphabet_size =
    if alphabet_size = 0 then 0
    else
      let w = ref 0 in
      let a = ref (alphabet_size - 1) in
      while !a > 0 do
        incr w;
        a := !a lsr 1
      done;
      Format.printf "width=%d" !w;
      !w

  let empty ~alphabet_size =
    {
      topol = 0L;
      labels = 0L;
      topol_length = 0;
      labels_length = 0;
      alphabet_size;
      width = width ~alphabet_size;
    }

  let open_node tree label =
    if
      tree.topol_length + 2 > max_width
      || tree.width * (tree.labels_length + 1) > max_width
    then None
    else
      let topol =
        Int64.logor tree.topol (Int64.shift_left 1L tree.topol_length)
      in
      let labels =
        let label64 = Int64.of_int label in
        let shifted =
          Int64.shift_left label64 (tree.width * tree.labels_length)
        in
        Int64.logor tree.labels shifted
      in
      Some
        {
          tree with
          topol;
          labels;
          topol_length = tree.topol_length + 1;
          labels_length = tree.labels_length + 1;
        }

  let close_node tree = { tree with topol_length = tree.topol_length + 1 }

  let finish_tree tree ~check =
    if not check then tree
    else
      let excess = ref 0 in
      let preorder = ref 0 in
      for i = 0 to tree.topol_length - 1 do
        if Int64.logand (Int64.shift_right_logical tree.topol i) 1L <> 0L then (
          incr excess;
          let mask = Int64.of_int ((1 lsl tree.width) - 1) in
          let shift = tree.width * !preorder in
          let lab =
            Int64.logand (Int64.shift_right_logical tree.labels shift) mask
          in
          if
            lab <> 0L && Int64.compare lab (Int64.of_int tree.alphabet_size) > 0
          then failwith (Printf.sprintf "Label %Lx exceeds sigma" lab);
          incr preorder)
        else (
          decr excess;
          if !excess < 0 then
            failwith (Printf.sprintf "Excess negative at position %d" i))
      done;
      if !excess <> 0 then failwith "Unbalanced parentheses";
      tree

  let pp ?fmt_label:(pp = fun fmt -> Format.fprintf fmt "%d") fmt tree =
    let tptr = ref 0 in
    let tlab = ref 0 in
    let excess = ref 0 in

    if tree.topol = 0L then Format.fprintf fmt "(null)"
    else
      while !tptr = 0 || !excess > 0 do
        let current_bit =
          Int64.logand (Int64.shift_right_logical tree.topol !tptr) 1L
        in
        if current_bit = 1L then (
          (* Open node *)
          let mask = Int64.of_int ((1 lsl tree.width) - 1) in
          let shift = tree.width * !tlab in
          let label_val =
            Int64.to_int
              (Int64.logand (Int64.shift_right_logical tree.labels shift) mask)
          in

          pp Format.std_formatter label_val;

          incr tptr;
          incr tlab;
          incr excess;

          let next_bit =
            Int64.logand (Int64.shift_right_logical tree.topol !tptr) 1L
          in
          if next_bit = 1L then Format.fprintf fmt "@[(";
          ())
        else (
          (* Close node *)
          incr tptr;
          decr excess;
          if !excess > 0 then
            let next_bit =
              Int64.logand (Int64.shift_right_logical tree.topol !tptr) 1L
            in
            if next_bit = 1L then Format.fprintf fmt ",@;<1 2>"
            else Format.fprintf fmt ")@]")
      done
end

module Tree = struct
  type t = Node of int * t list

  let ( let* ) x f = Option.bind x f

  let to_succint_tree ~alphabet_size tree =
    let empty = SuccinctTree.empty ~alphabet_size in
    let rec go (t : t) (acc : 'a option) =
      let* acc = acc in
      match t with
      | Node (id, []) ->
          let* node = SuccinctTree.open_node acc id in
          Option.some @@ SuccinctTree.close_node node
      | Node (n, l) ->
          let st1 = SuccinctTree.open_node acc n in
          let* st2 = List.fold_left (fun acc child -> go child acc) st1 l in
          let st3 = SuccinctTree.close_node st2 in
          Option.some st3
    in
    let* succinct_tree = go tree @@ Option.some empty in
    Option.some @@ SuccinctTree.finish_tree ~check:true succinct_tree
end

open Syntax
open Closure
open Lower
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`

| Value   | Block                             |
|---------|----------------------------------:|
| Tuple   | [ ev ; len ; vals ... ]           |
| Closure | [ ev ; fptr ; captured vals ... ] |
| Vec     | [ ev ; vals ... ]                 |
| Variant | [ ev ; id ; vals ... ]            |

### 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`)

### Evidence

Evidence is implemented with succint trees. This is a memory efficient representation of trees.

Evidence takes two 64 bit words. The first one for the topology, the second one for the labels.

*)

(* number of 8 byte words used by evidence *)
let evidence_word_num = 2
let alphabet_size = 16

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

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

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

(* variant ids must be larger or equal than this number *)
let variant_id_base = 7

let variant_id ty =
  let id = Type.variant_id ty + variant_id_base in
  if id >= alphabet_size then
    raise (Invalid_argument "variant id exceeds the alphabet size");
  id

let id_of_ty (ty : Type.t) =
  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 "should never happen"

let children_of_ty (ty : Type.t) =
  match ty with
  | Type.Any -> []
  | Type.Int -> []
  | Type.Bool -> []
  | Type.Float -> []
  | Type.Variant _ | Type.Rec (_, _) -> []
  | Type.Arrow (ty_args, ty_ret) -> List.append ty_args [ ty_ret ]
  | Type.Product tys -> tys
  | Type.Vec ty -> [ ty ]
  | Type.Var _ -> failwith "should never happen"

let rec tree_of_ty (ty : Type.t) =
  let id = id_of_ty ty in
  let children = children_of_ty ty in
  Tree.Node (id, List.map tree_of_ty children)

let succinct_tree_of_ty_opt ty =
  let tree = tree_of_ty ty in
  Tree.to_succint_tree ~alphabet_size tree

let succinct_tree_of_ty ty =
  match succinct_tree_of_ty_opt ty with
  | None ->
      Format.eprintf "Tree too large to convert into succinct tree: %a" Type.pp
        ty;
      failwith "tree too large"
  | Some t -> t

let unexpected_variant_case_error = 15L
let const_false = 0b010L
let const_true = 0b110L

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

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), evidence_word_num))

let lower_ev (ev : Type.t) =
  let succinct_tree = succinct_tree_of_ty ev in
  ( Cmm.Syntax.Int (Int64, succinct_tree.topol),
    Cmm.Syntax.Int (Int64, succinct_tree.labels) )

module GradualMemrep :
  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 (c : Const.t) : t =
    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)

  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 ev =
    let open Cmm.Syntax in
    let num_captured_vars =
      Int (Int64, Int64.of_int (List.length captured_vars))
    in
    let bind x f b = Let (x, f, b) in
    let store t t_value idx = Store (t_value, Gep (TInt Int64, t, idx)) in
    let topol, labels = lower_ev ev in
    bind "closure_ptr"
      (ApplyExt
         ( "allocate_closure",
           [ TPtr; TInt Int64 ],
           TPtr,
           [ Func fname; num_captured_vars ] ))
    @@ Seq
         [
           store (Var "closure_ptr") topol 0;
           store (Var "closure_ptr") labels 1;
           Ptrtoint (Var "closure_ptr", TInt Int64);
         ]

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

  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, Load (TPtr, Gep (TPtr, Var "self", 3)), idx),
               Let (name, Load (TInt Int64, Var "ptr"), acc) ))
         body

  and tuple : 'a atom list -> Cmm.Syntax.expr =
   fun atoms ->
    let open Cmm.Syntax in
    let len = List.length atoms in
    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
                ( "tuple",
                  ApplyExt
                    ( "make_tuple",
                      [ TInt Int64; TPtr ],
                      TInt Int64,
                      [ length; Var "vals" ] ),
                  Ptrtoint (Var "tuple", 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 = 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_ptr",
        allocate_variant,
        Let
          ( "variant",
            Gep (TArray (TInt Int64, 0), Var "variant_ptr", 0),
            Seq (fill_variant @ [ Ptrtoint (Var "variant_ptr", TInt Int64) ]) )
      )
end

module GradualElim : 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 -> tag_int @@ Prim1 (Neg, untag_int 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 to_i64 t =
      Cmm.Syntax.Prim2 (And, Intcast (t, TInt Int64), Int (Int64, 1L))
    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 (to_i64 (Prim2 (Eq, v1, v2)))
    | Prim.Less -> tag_bool (to_i64 (Prim2 (Lt, v1, v2)))
    | Prim.LessEq -> tag_bool (to_i64 (Prim2 (Le, v1, v2)))
    | Prim.Greater -> tag_bool (to_i64 (Prim2 (Gt, v1, v2)))
    | Prim.GreaterEq -> tag_bool (to_i64 (Prim2 (Ge, v1, v2)))
    | Prim.FPlus ->
        box_float @@ Prim2 (Cmm.Syntax.FPlus, unbox_float v1, unbox_float v2)
    | Prim.FMinus ->
        box_float @@ Prim2 (Cmm.Syntax.FMinus, unbox_float v1, unbox_float v2)
    | Prim.FTimes ->
        box_float @@ Prim2 (Cmm.Syntax.FTimes, unbox_float v1, unbox_float v2)
    | Prim.FDiv ->
        box_float @@ Prim2 (Cmm.Syntax.FDiv, unbox_float v1, unbox_float 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 ] )
    | 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 : t -> int -> t =
   fun ptr idx ->
    Cmm.Syntax.ApplyExt
      ( "tuple_get",
        [ TInt Int64; TInt Int64 ],
        TInt Int64,
        [ ptr; Int (Int64, Int64.of_int idx) ] )

  let ascribe : t -> Type.t -> t =
   fun value ty ->
    let t_of_int n = Cmm.Syntax.Int (Int64, Int64.of_int n) 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 _ ->
        let topol, labels = lower_ev ty in
        Cmm.Syntax.ApplyExt
          ( "ascribe_closure",
            [ TInt Int64; TInt Int64; TInt Int64 ],
            TInt Int64,
            [ value; labels; topol ] )
    | Vec ty_inner ->
        let topol, labels = lower_ev ty_inner in
        Cmm.Syntax.ApplyExt
          ( "ascribe_vec",
            [ TInt Int64; TInt Int64; TInt Int64 ],
            TInt Int64,
            [ value; labels; topol ] )
    | Product _ ->
        let topol, labels = lower_ev ty in
        Cmm.Syntax.ApplyExt
          ( "ascribe_tuple",
            [ TInt Int64; TInt Int64; TInt Int64 ],
            TInt Int64,
            [ value; labels; topol ] )
    | Float | Variant _ | Rec (_, _) ->
        let id = id_of_ty ty in
        Cmm.Syntax.ApplyExt
          ( "ascribe_imm_evidence",
            [ TInt Int64; TInt Int64 ],
            TInt Int64,
            [ value; t_of_int id ] )
    | Var _ -> failwith "this case should be avoided by the normalization phase"

  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 bind_list = List.fold_right2 bind in
    let seq l b = Seq (l @ [ b ]) in
    let load t idx =
      Load (TInt Int64, Gep (TInt Int64, Inttoptr (t, TPtr), idx))
    in
    let i64_t = TInt Int64 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 ascribe_args =
      List.map
        (fun arg ->
          ApplyExt
            ("ascribe_arg", [ i64_t; TPtr ], i64_t, [ arg; Var "iterator" ]))
        args
      @ [
          ApplyExt
            ("assert_single_type_remaining", [ TPtr ], TVoid, [ Var "iterator" ]);
        ]
    in
    bind "closure_ptr" (to_ptr closure)
    @@ bind "f_ptr" (to_ptr (load (Var "closure_ptr") evidence_word_num))
    @@ bind "iterator" (stack_alloc (TArray (TInt Int64, 2)))
    @@ seq
         [
           ApplyExt ("assert_arrow", [ TPtr ], TVoid, [ Var "closure_ptr" ]);
           ApplyExt
             ( "typeTree_children_init",
               [ TPtr; TPtr ],
               TVoid,
               [ Var "iterator"; Var "closure_ptr" ] );
         ]
    @@ bind_list arg_names args
    @@ bind "ret"
         (Seq
            (ascribe_args
            @ [ Apply (func, (Var "closure_ptr", TPtr) :: args', i64_t) ]))
    @@ bind "ret_asc"
         (ApplyExt
            ( "ascribe_arg",
              [ i64_t; TPtr ],
              i64_t,
              [ Var "ret"; Var "iterator" ] ))
    @@ Var "ret_asc"

  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_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_ptr",
        Inttoptr (variant_ptr, TPtr),
        Let
          ( "case_id",
            Load
              ( TInt Int64,
                Gep (TInt Int64, Var "variant_ptr", evidence_word_num) ),
            Switch (Var "case_id", cases (Var "variant_ptr"), default) ) )
end

module LowerStandardMainFunction = MakeLower (GradualMemrep) (GradualElim)

module Lower = struct
  include LowerStandardMainFunction

  let codegen_program 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
              [
                ApplyExt ("preamble", [], TVoid, []);
                codegen_expr program.expr;
                Int (Int32, 0L);
              ];
        }
    in
    funcs @ [ func ]
end
