open Common

let error msg = Error.error (Elaborate msg)

module type Evidence = sig
  type g
  type t

  val interior : g -> g -> t
  val meet : g -> g -> g
  val precision : g -> g -> bool option
  val consistent_transitivity : t -> t -> t option
  val is_fully_precise : t -> bool
  val invert : t -> t
  val pp : Format.formatter -> t -> unit
  val can_be_optimized : t -> bool
end

module type TypeMatching = sig
  type t

  val cod : t -> t option
  val dom : t -> int -> t option
  val vec : t -> t option
  val proj : t -> int -> t option
end

module TypeMatching = struct
  type t = Type.t

  let cod (t : t) : t option =
    match t with Any -> Some Any | Arrow (_, ty) -> Some ty | _ -> None

  let dom (t : t) (i : int) : t option =
    match t with
    | Any -> Some Type.Any
    | Arrow (tys, _) -> List.nth_opt tys i
    | _ -> None

  let vec (t : t) : t option =
    match t with Any -> Some Any | Vec ty -> Some ty | _ -> None

  let proj (t : t) (i : int) : t option =
    match t with
    | Any -> Some Type.Any
    | Product tys -> List.nth_opt tys i
    | _ -> None
end

module _ : TypeMatching = TypeMatching

module TypeEvidence = struct
  type g = Type.t
  type t = Type.t

  let pp = Type.pp

  let meet t1 t2 =
    let rec meet t1 t2 env =
      match (t1, t2) with
      | Type.Var x, t ->
          if List.mem_assoc x env then meet (List.assoc x env) t env
          else raise (error (Format.asprintf "unbound type variable '%s'" x))
      | t, Type.Var x ->
          if List.mem_assoc x env then meet t (List.assoc x env) env
          else raise (error (Format.asprintf "unbound type variable '%s'" x))
      | Type.Any, t | t, Type.Any -> t
      | Type.Bool, Type.Bool -> Type.Bool
      | Type.Int, Type.Int -> Type.Int
      | Type.Float, Type.Float -> Type.Float
      | Type.Arrow (tys, ty), Type.Arrow (tys', ty') ->
          if List.length tys <> List.length tys' then
            raise (error "Arrow types of different lengths do not match");
          Type.Arrow
            (List.map2 (fun t1 t2 -> meet t1 t2 env) tys tys', meet ty ty' env)
      | Type.Product ts, Type.Product ts' ->
          if List.length ts <> List.length ts' then
            raise (error "Product types of different lengths do not match");
          Type.Product (List.map2 (fun t1 t2 -> meet t1 t2 env) ts ts')
      | Type.Variant ts, Type.Variant ts' ->
          if List.length ts <> List.length ts' then
            raise (error "Variant types of different lengths do not match");
          if ts = ts' then t1
          else
            raise
              (error
                 (Format.asprintf "Variant types '%a' and '%a' do not match"
                    Type.pp t1 Type.pp t2))
      | Type.Vec t1, Type.Vec t2 -> Type.Vec (meet t1 t2 env)
      | Type.Rec (x, t), Type.Rec (x', t') ->
          Type.Rec (x, meet t t' ((x, Type.Var x') :: env))
      | Type.Bool, _
      | Type.Int, _
      | Type.Float, _
      | Type.Arrow _, _
      | Type.Product _, _
      | Type.Vec _, _
      | Type.Rec _, _
      | Type.Variant _, _ ->
          raise
            (error
               (Format.asprintf "meet is undefined for types '%a' and '%a'"
                  Type.pp t1 Type.pp t2))
    in
    meet t1 t2 []

  let interior t1 t2 =
    let m = meet t1 t2 in
    m

  let precision ty1 ty2 =
    try
      let ty_meet = meet ty1 ty2 in
      Some (ty1 = ty_meet)
    with _ -> None

  let consistent_transitivity ty1 ty2 = try Some (meet ty1 ty2) with _ -> None

  let rec is_fully_precise = function
    | Type.Any -> false
    | Type.Int | Type.Bool | Type.Float | Type.Rec (_, _) | Type.Variant _ ->
        true
    | Type.Arrow (tys, ty) ->
        List.for_all is_fully_precise tys && is_fully_precise ty
    | Type.Product tys -> List.for_all is_fully_precise tys
    | Type.Vec ty -> is_fully_precise ty
    | Type.Var _ -> failwith "should be impossible"

  let invert t = t
  let can_be_optimized _ev = true
end

module RuntimeType = struct
  type t =
    | Any
    | Int
    | Bool
    | Float (* FF *)
    | BoxedFloat (* BB *)
    | FloatOfBox (* BF *)
    | BoxOfFloat (* FB *)
    | Arrow of t list * t
    | Product of t list
    | Vec of t
    | Variant of case list
    | Rec of string * t
    | Var of string

  and case = Case of string * t list

  let rec pp fmt = function
    | Int -> Format.fprintf fmt "int"
    | Bool -> Format.fprintf fmt "bool"
    | Float -> Format.fprintf fmt "float"
    | FloatOfBox -> Format.fprintf fmt "box->float"
    | BoxOfFloat -> Format.fprintf fmt "float->box"
    | BoxedFloat -> Format.fprintf fmt "boxed_float"
    | Any -> Format.fprintf fmt "?"
    | Arrow (tys, ty) ->
        Format.fprintf fmt "(%a -> %a)" (Pp.pp_list ~sep:" -> " pp) tys pp ty
    | Product l -> Format.fprintf fmt "(%a)" (Pp.pp_list ~sep:" * " pp) l
    | Vec t -> Format.fprintf fmt "vec[%a]" pp t
    | Variant cases ->
        Format.fprintf fmt "(%a)" (Pp.pp_list ~sep:" | " pp_case) cases
    | Rec (type_name, _) -> Format.fprintf fmt "%s" type_name
    | Var type_name -> Format.fprintf fmt "%s" type_name

  and pp_case fmt = function
    | Case (name, tys) ->
        if tys = [] then Format.fprintf fmt "%s" name
        else Format.fprintf fmt "%s of %a" name (Pp.pp_list ~sep:" * " pp) tys

  let variant_id =
    let hashtable = Hashtbl.create 10 in
    fun ty ->
      match ty with
      | Rec (_, Variant _) ->
          if Hashtbl.mem hashtable ty then Hashtbl.find hashtable ty
          else
            let id = Hashtbl.length hashtable in
            Hashtbl.add hashtable ty id;
            id
      | _ -> failwith "tried to assign variant id to type that is not a variant"

  let case_id casename ty =
    let cases =
      match ty with
      | Rec (_, Variant cases) -> cases
      | _ -> failwith "tried to obtain case id from type that is not a variant"
    in
    let i =
      cases
      |> List.mapi (fun i case -> (i, case))
      |> List.find (fun (_, Case (name, _)) -> name = casename)
      |> fst
    in
    i
end

module RuntimeTypeMatching = struct
  type t = RuntimeType.t

  let cod (t : t) : t option =
    match t with Any -> Some Any | Arrow (_, ty) -> Some ty | _ -> None

  let dom (t : t) (i : int) : t option =
    match t with
    | Any -> Some Any
    | Arrow (tys, _) -> List.nth_opt tys i
    | _ -> None

  let vec (t : t) : t option =
    match t with Any -> Some Any | Vec ty -> Some ty | _ -> None

  let proj (t : t) (i : int) : t option =
    match t with
    | Any -> Some Any
    | Product tys -> List.nth_opt tys i
    | _ -> None
end

module _ : TypeMatching = RuntimeTypeMatching

module RuntimeTypeEvidence = struct
  type t = RuntimeType.t
  type g = Type.t

  module T = Type
  module R = RuntimeType

  let pp = RuntimeType.pp

  let interior t1 t2 =
    let rec interior (t1 : g) (t2 : g) env : t =
      match (t1, t2) with
      | T.Float, T.Float -> Float
      | T.Float, T.Any -> BoxOfFloat
      | T.Any, T.Float -> FloatOfBox
      | T.Any, T.Any -> Any
      | T.Bool, T.Bool | T.Bool, T.Any | T.Any, T.Bool -> Bool
      | T.Int, T.Int | T.Int, T.Any | T.Any, T.Int -> Int
      | T.Arrow (tys, ty), T.Arrow (tys', ty') ->
          if List.length tys <> List.length tys' then
            raise (error "Arrow types of different lengths do not match");
          Arrow
            ( List.map2 (fun t1 t2 -> interior t1 t2 env) tys' tys,
              interior ty ty' env )
      | T.Arrow (tys, ty), T.Any ->
          Arrow
            ( List.map (fun ty -> interior T.Any ty env) tys,
              interior ty T.Any env )
      | T.Any, T.Arrow (tys, ty) ->
          Arrow
            ( List.map (fun ty -> interior ty T.Any env) tys,
              interior T.Any ty env )
      | T.Product ts, T.Product ts' ->
          if List.length ts <> List.length ts' then
            raise (error "Product types of different lengths do not match");
          Product (List.map2 (fun t1 t2 -> interior t1 t2 env) ts ts')
      | T.Product ts, T.Any ->
          Product (List.map (fun ty -> interior ty Any env) ts)
      | T.Any, T.Product ts ->
          Product (List.map (fun ty -> interior Any ty env) ts)
      | T.Vec t1, T.Vec t2 -> Vec (interior t1 t2 env)
      | T.Vec ty, T.Any -> Vec (interior ty Any env)
      | T.Any, T.Vec ty -> Vec (interior Any ty env)
      | T.Variant cases, T.Variant cases' ->
          if List.length cases <> List.length cases' then
            raise (error "Variant types of different lengths do not match");
          if cases <> cases' then
            raise
              (error
                 (Format.asprintf "Variant types '%a' and '%a' do not match"
                    T.pp t1 T.pp t2));
          Variant
            (List.map
               (fun (T.Case (x, tys)) ->
                 R.Case
                   (x, List.map2 (fun ty1 ty2 -> interior ty1 ty2 env) tys tys))
               cases)
      | T.Variant cases, T.Any ->
          Variant
            (List.map
               (fun (T.Case (x, tys)) ->
                 R.Case (x, List.map (fun ty -> interior ty Any env) tys))
               cases)
      | T.Any, T.Variant cases ->
          Variant
            (List.map
               (fun (T.Case (x, tys)) ->
                 R.Case (x, List.map (fun ty -> interior Any ty env) tys))
               cases)
      | T.Rec (x, t), T.Rec (x', t') ->
          Rec (x, interior t t' ((x, T.Var x) :: (x', T.Var x) :: env))
      | T.Rec (x, t), T.Any -> Rec (x, interior t Any ((x, T.Var x) :: env))
      | T.Any, T.Rec (x, t) -> Rec (x, interior Any t ((x, T.Var x) :: env))
      | T.Var x, T.Var x' when x = x' -> Var x
      | T.Var x, t ->
          if List.mem_assoc x env then
            let t' = List.assoc x env in
            if t' = T.Var x then Var x else interior t' t env
          else raise (error (Format.asprintf "unbound type variable '%s'" x))
      | t, T.Var x ->
          if List.mem_assoc x env then
            let t' = List.assoc x env in
            if t' = T.Var x then Var x else interior t t' env
          else raise (error (Format.asprintf "unbound type variable '%s'" x))
      | T.Bool, _
      | T.Int, _
      | T.Float, _
      | T.Arrow _, _
      | T.Product _, _
      | T.Vec _, _
      | T.Rec _, _
      | T.Variant _, _ ->
          raise
            (error
               (Format.asprintf "meet is undefined for types '%a' and '%a'" T.pp
                  t1 T.pp t2))
    in
    interior t1 t2 []

  let meet = TypeEvidence.meet

  let precision ty1 ty2 =
    try
      let ty_meet = meet ty1 ty2 in
      Some (ty1 = ty_meet)
    with _ -> None

  let rec consistent_transitivity (ev1 : t) (ev2 : t) : t option =
    let ( let* ) x f = Option.bind x f in
    let option_map l =
      if List.for_all Option.is_some l then Some (List.map Option.get l)
      else None
    in
    match (ev1, ev2) with
    | Int, Int -> Some Int
    | Bool, Bool -> Some Bool
    | Float, Float -> Some Float (* FF o FF = FF*)
    | Float, BoxOfFloat -> Some BoxOfFloat (* FF o FB = FB*)
    | Float, FloatOfBox -> None (* FF o BF = error *)
    | Float, BoxedFloat -> None (* FF o BB = error *)
    | BoxOfFloat, Float -> None (* FB o FF = error *)
    | BoxOfFloat, BoxOfFloat -> None (* FB o FB = error *)
    | BoxOfFloat, FloatOfBox -> Some Float (* FB o BF = FF *)
    | BoxOfFloat, BoxedFloat -> Some BoxOfFloat (* FB o BB = FB *)
    | BoxedFloat, Float -> None (* BB o FF = error *)
    | BoxedFloat, BoxOfFloat -> None (* BB o FB = error *)
    | BoxedFloat, FloatOfBox -> Some FloatOfBox (* BB o BF = FF *)
    | BoxedFloat, BoxedFloat -> Some BoxedFloat (* BB o BB = BB *)
    | FloatOfBox, Float -> Some FloatOfBox (* BF o FF = FF *)
    | FloatOfBox, BoxOfFloat -> Some BoxedFloat (* BF o FB = BB *)
    | FloatOfBox, FloatOfBox -> None (* BF o BF = error *)
    | FloatOfBox, BoxedFloat -> None (* BF o BB = error *)
    | Any, Any -> Some Any
    | Any, Int -> Some Int
    | Int, Any -> Some Int
    | Any, Bool -> Some Bool
    | Bool, Any -> Some Bool
    | Any, Arrow (tys, _) ->
        let tys' = List.map (fun _ -> R.Any) tys in
        let ty' = R.Any in
        consistent_transitivity (R.Arrow (tys', ty')) ev2
    | Arrow (tys, _), Any ->
        let tys' = List.map (fun _ -> R.Any) tys in
        let ty' = R.Any in
        consistent_transitivity ev1 (R.Arrow (tys', ty'))
    | Any, Product tys ->
        let tys' = List.map (fun _ -> R.Any) tys in
        consistent_transitivity (R.Product tys') ev2
    | Product tys, Any ->
        let tys' = List.map (fun _ -> R.Any) tys in
        consistent_transitivity ev1 (R.Product tys')
    | Any, Vec _ ->
        let ty' = R.Any in
        consistent_transitivity (R.Vec ty') ev2
    | Vec _, Any ->
        let ty' = R.Any in
        consistent_transitivity ev1 (R.Vec ty')
    | Any, Rec (x, _) -> consistent_transitivity (R.Rec (x, R.Any)) ev2
    | Rec (x, _), R.Any -> consistent_transitivity ev1 (R.Rec (x, R.Any))
    | Any, Variant cases -> Some (Variant cases)
    | Variant cases, Any -> Some (Variant cases)
    | Any, Var x -> Some (Var x)
    | Var x, Any -> Some (Var x)
    | Arrow (tys, ty), Arrow (tys', ty') ->
        if List.length tys = List.length tys' then
          let* tys = List.map2 consistent_transitivity tys' tys |> option_map in
          let* ty = consistent_transitivity ty ty' in
          Some (R.Arrow (tys, ty))
        else None
    | Product tys, Product tys' ->
        if List.length tys = List.length tys' then
          let* tys = List.map2 consistent_transitivity tys tys' |> option_map in
          Some (R.Product tys)
        else None
    | Vec ty, Vec ty' ->
        let* ty = consistent_transitivity ty ty' in
        Some (R.Vec ty)
    | Rec (x, t), Rec (x', t') ->
        if x = x' then
          let* t_combined = consistent_transitivity t t' in
          Some (R.Rec (x, t_combined))
        else None
    | Variant cases, Variant cases' ->
        if cases = cases' then Some (Variant cases) else None
    | Var x, Var y -> if x = y then None else Some (Var x)
    | _ -> None

  let rec is_fully_precise (ev : t) =
    match ev with
    | Any | BoxOfFloat | FloatOfBox | BoxedFloat -> false
    | Int | Bool | Float | Rec (_, _) | Variant _ -> true
    | Arrow (tys, ty) ->
        List.for_all is_fully_precise tys && is_fully_precise ty
    | Product tys -> List.for_all is_fully_precise tys
    | Vec ty -> is_fully_precise ty
    | Var _ -> failwith "should be impossible"

  let rec invert (ev : t) =
    match ev with
    | Any | Int | Bool | Float | BoxedFloat -> ev
    | FloatOfBox -> BoxOfFloat
    | BoxOfFloat -> FloatOfBox
    | Arrow (ts, t) -> Arrow (List.map invert ts, invert t)
    | Product ts -> Product (List.map invert ts)
    | Vec t -> Vec (invert t)
    | Rec (_, _) | Variant _ | Var _ -> ev

  let rec can_be_optimized (ev : t) =
    match ev with
    | Float | BoxedFloat | FloatOfBox | BoxOfFloat | Any -> false
    | Int | Bool -> true
    | Arrow (tys, ty) ->
        List.for_all can_be_optimized tys && can_be_optimized ty
    | Product tys -> List.for_all can_be_optimized tys
    | Vec ty -> can_be_optimized ty
    | Rec (_, ty) -> can_be_optimized ty
    | Variant _ -> true
    | Var _ -> true
end

module type Elaborate = sig
  type ev

  val elaborate : Typed_ast.e -> ev Ast.t
end

module Elaborate (Ev : Evidence with type g = Type.t) = struct
  module Env = Map.Make (String)

  let is_value (t : Ev.t Ast.t) : bool =
    match t.desc with
    | Const _ | Lam _ | Tuple _ | Variant _ | Vect _ -> true
    | _ -> false

  let rec elaborate (program : Typed_ast.program) : Ev.t Ast.t =
    let open Typed_ast in
    let env =
      List.fold_left
        (fun env (DType (name, ty)) -> Env.add name (Ev.interior ty ty, ty) env)
        Env.empty program.decls
    in
    let expr, _ = elaborate' program.expr env in
    expr

  and elaborate' (t : Typed_ast.e) (env : (Ev.t * Type.t) Env.t) :
      Ev.t Ast.t * Type.t =
    let err msg = Error.error_with_loc (Elaborate msg) t.loc in
    let interior ty1 ty2 =
      try Ev.interior ty1 ty2
      with Error.Error (_, err) -> raise (Error.error_with_loc err t.loc)
    in
    let make desc meta = Ast.{ desc; meta } in
    match t.desc with
    | Const c ->
        let ty = Const.type_of c in
        let ev = interior ty ty in
        (make (Const c) ev, ty)
    | Var x ->
        let ev, ty =
          match Env.find_opt x env with
          | Some (ev, ty) -> (ev, ty)
          | None -> raise (err (Format.asprintf "unbound variable '%s'" x))
        in
        (make (Var x) ev, ty)
    | Lam (args, body) ->
        let env' =
          List.fold_left
            (fun env (x, ty) -> Env.add x (interior ty ty, ty) env)
            env args
        in
        let body, ty' = elaborate' body env' in
        let xs, tys = List.split args in
        let ty_arrow = Type.Arrow (tys, ty') in
        let ev = interior ty_arrow ty_arrow in
        (make (Lam (xs, body)) ev, ty_arrow)
    | App (f, args) ->
        let f, ty_f = elaborate' f env in
        let args, ty_args =
          List.split (List.map (fun arg -> elaborate' arg env) args)
        in
        let arity = List.length args in
        let ty_arg i =
          match TypeMatching.dom ty_f i with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf
                      "expected a function with arity %d, but got %a" arity
                      Type.pp ty_f))
        in
        let ty_args' = List.init arity (fun i -> ty_arg i) in
        let ty_ret' =
          match TypeMatching.cod ty_f with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf "expected a function type, but got %a"
                      Type.pp ty_f))
        in
        let ev_f = interior ty_f (Type.Arrow (ty_args', ty_ret')) in
        let ev_args = List.map2 interior ty_args ty_args' in
        let f_asc = ascribe f ev_f (Type.Arrow (ty_args', ty_ret')) in
        let args_asc =
          List.map2
            (fun (arg, ev) ty -> ascribe arg ev ty)
            (List.combine args ev_args)
            ty_args'
        in
        (make (App (f_asc, args_asc)) (interior ty_ret' ty_ret'), ty_ret')
    | Prim1 (op1, arg) ->
        let arg, ty_arg = elaborate' arg env in
        let ty1', ty2' = Type.ty_of_op1 op1 in
        let ev = interior ty_arg ty1' in
        let arg_asc = ascribe arg ev ty1' in
        (make (Prim1 (op1, arg_asc)) (interior ty2' ty2'), ty2')
    | Prim2 (op2, arg1, arg2) ->
        let arg1, ty_arg1 = elaborate' arg1 env in
        let arg2, ty_arg2 = elaborate' arg2 env in
        let ty1', ty2', ty3' = Type.ty_of_op2 op2 in
        let ev1 = interior ty_arg1 ty1' in
        let ev2 = interior ty_arg2 ty2' in
        let arg1_asc = ascribe arg1 ev1 ty1' in
        let arg2_asc = ascribe arg2 ev2 ty2' in
        (make (Prim2 (op2, arg1_asc, arg2_asc)) (interior ty3' ty3'), ty3')
    | If (arg_cond, arg_true, arg_false) ->
        let arg_cond, ty_arg_cond = elaborate' arg_cond env in
        let arg_true, ty_arg_true = elaborate' arg_true env in
        let arg_false, ty_arg_false = elaborate' arg_false env in
        let ev_cond = interior ty_arg_cond Type.Bool in
        let ty_args = Ev.meet ty_arg_true ty_arg_false in
        let ev_arg_true = interior ty_arg_true ty_args in
        let ev_arg_false = interior ty_arg_false ty_args in
        let arg_cond_asc = ascribe arg_cond ev_cond ty_arg_cond in
        let arg_true_asc = ascribe arg_true ev_arg_true ty_arg_true in
        let arg_false_asc = ascribe arg_false ev_arg_false ty_arg_false in
        ( make
            (If (arg_cond_asc, arg_true_asc, arg_false_asc))
            (interior ty_arg_true ty_arg_true),
          ty_args )
    | Let (binding, body) ->
        let term, ty_term = elaborate' binding.e env in
        let ev_term = interior ty_term binding.ty in
        let term_asc = ascribe term ev_term binding.ty in
        let body, ty_body =
          elaborate' body (Env.add binding.name (ev_term, binding.ty) env)
        in
        let binding' = Ast.{ term = term_asc; name = binding.name } in
        (make (Let (binding', body)) body.meta, ty_body)
    | Letrec (bindings, body) ->
        let env' =
          List.fold_left
            (fun env (binding : Typed_ast.binding) ->
              let ev = interior binding.ty binding.ty in
              Env.add binding.name (ev, binding.ty) env)
            env bindings
        in
        let bindings' =
          bindings
          |> List.map (fun (binding : Typed_ast.binding) ->
                 let term, ty_term = elaborate' binding.e env' in
                 let term_asc =
                   ascribe term (interior ty_term binding.ty) binding.ty
                 in
                 Ast.{ term = term_asc; name = binding.name })
        in
        let body, ty_body = elaborate' body env' in
        (make (Letrec (bindings', body)) body.meta, ty_body)
    | Tuple args ->
        let args, ty_args =
          List.split (List.map (fun arg -> elaborate' arg env) args)
        in
        let ty_tup = Type.Product ty_args in
        let args_asc =
          List.combine args ty_args
          |> List.map (fun (arg, ty_arg) ->
                 let arg_ev = interior ty_arg ty_arg in
                 ascribe arg arg_ev ty_arg)
        in
        let ev = interior ty_tup ty_tup in
        (make (Tuple args_asc) ev, ty_tup)
    | Get (arg, i) ->
        let arg, ty_arg = elaborate' arg env in
        let ty =
          match TypeMatching.proj ty_arg i with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf
                      "expected a product type with arity of at least %d, but \
                       got %a"
                      i Type.pp ty_arg))
        in
        let ev = interior ty ty in
        let arg_ev = interior ty_arg ty_arg in
        let arg_asc = ascribe arg arg_ev ty_arg in
        (make (Get (arg_asc, i)) ev, ty)
    | Vect (arg1, arg2) ->
        let arg1, ty_arg1 = elaborate' arg1 env in
        let arg2, ty_arg2 = elaborate' arg2 env in
        let ty_vec = Type.Vec ty_arg2 in
        let ev = interior ty_vec ty_vec in
        let arg1_ev = interior ty_arg1 ty_arg1 in
        let arg1_asc = ascribe arg1 arg1_ev Type.Int in
        let arg2_ev = interior ty_arg2 ty_arg2 in
        let arg2_asc = ascribe arg2 arg2_ev ty_arg2 in
        (make (Vect (arg1_asc, arg2_asc)) ev, ty_vec)
    | VectGet (arg1, arg2) ->
        let arg1, ty_arg1 = elaborate' arg1 env in
        let arg2, ty_arg2 = elaborate' arg2 env in
        let ty =
          match TypeMatching.vec ty_arg1 with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf "expected a vector type, but got %a" Type.pp
                      ty_arg1))
        in
        let ev = interior ty ty in
        let arg1_ev = interior ty_arg1 (Type.Vec ty) in
        let arg1_asc = ascribe arg1 arg1_ev (Type.Vec ty) in
        let arg2_ev = interior ty_arg2 ty_arg2 in
        let arg2_asc = ascribe arg2 arg2_ev Type.Int in
        (make (VectGet (arg1_asc, arg2_asc)) ev, ty)
    | VectSet (arg1, arg2, arg3) ->
        let arg1, ty_arg1 = elaborate' arg1 env in
        let arg2, ty_arg2 = elaborate' arg2 env in
        let arg3, ty_arg3 = elaborate' arg3 env in
        let ty_inner =
          match TypeMatching.vec ty_arg1 with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf "expected a vector type, but got %a" Type.pp
                      ty_arg1))
        in
        let ty = Type.Vec ty_inner in
        let ev = interior ty (Type.Vec ty_arg3) in
        let arg1_ev = interior ty_arg1 ty in
        let arg1_asc = ascribe arg1 arg1_ev ty in
        let arg2_ev = interior ty_arg2 ty_arg2 in
        let arg2_asc = ascribe arg2 arg2_ev Type.Int in
        let arg3_ev = interior ty_arg3 ty_inner in
        let arg3_asc = ascribe arg3 arg3_ev ty_inner in
        (make (VectSet (arg1_asc, arg2_asc, arg3_asc)) ev, ty)
    | VectLength arg ->
        let arg, ty_arg = elaborate' arg env in
        let ty =
          match TypeMatching.vec ty_arg with
          | Some ty -> ty
          | None ->
              raise
                (err
                   (Format.asprintf "expected a vector type, but got %a" Type.pp
                      ty_arg))
        in
        let ev = interior Int Int in
        let arg_ev = interior ty_arg (Type.Vec ty) in
        let arg_asc = ascribe arg arg_ev (Type.Vec ty) in
        (make (VectLength arg_asc) ev, Type.Int)
    | Ascribe (arg, ty) ->
        let arg', ty_arg = elaborate' arg env in
        let ev = interior ty_arg ty in
        (ascribe arg' ev ty, ty)
    | For (x, arg1, arg2, arg3) ->
        let env' = Env.add x (interior Type.Int Type.Int, Type.Int) env in
        let arg1, ty_arg1 = elaborate' arg1 env in
        let arg2, ty_arg2 = elaborate' arg2 env in
        let arg3, _ = elaborate' arg3 env' in
        let ev1 = interior ty_arg1 Type.Int in
        let ev2 = interior ty_arg2 Type.Int in
        let arg1_asc = ascribe arg1 ev1 Type.Int in
        let arg2_asc = ascribe arg2 ev2 Type.Int in
        let ty = Type.Int in
        let ev = interior ty ty in
        (make (For (x, arg1_asc, arg2_asc, arg3)) ev, ty)
    | Variant (name, args, ty) ->
        let ty' =
          match ty with
          | Type.Rec (name, ty_inner) -> Type.subst_ty name ty ty_inner
          | _ ->
              raise
                (err
                   (Format.asprintf "expected rec variant type, but got '%a'"
                      Type.pp ty))
        in
        let cases =
          match ty' with
          | Type.Variant cases -> cases
          | _ ->
              raise
                (err
                   (Format.asprintf "expected variant type, but got '%a'"
                      Type.pp ty))
        in
        let case =
          cases |> List.find (fun (Type.Case (name', _)) -> name' = name)
        in
        let tys = match case with Type.Case (_, tys) -> tys in
        let args_asc =
          List.map2
            (fun arg ty ->
              let arg, ty_arg = elaborate' arg env in
              let ev = interior ty_arg ty in
              ascribe arg ev ty)
            args tys
        in
        let ev = interior ty ty in
        (make (Variant (name, args_asc, ty)) ev, ty)
    | Match (arg, patterns, ty) ->
        let ty' =
          match ty with
          | Type.Rec (name, ty_inner) -> Type.subst_ty name ty ty_inner
          | _ ->
              raise
                (err
                   (Format.asprintf "expected rec variant type, but got '%a'"
                      Type.pp ty))
        in
        let cases =
          match ty' with
          | Type.Variant cases -> cases
          | _ ->
              raise
                (err
                   (Format.asprintf "expected variant type, but got '%a'"
                      Type.pp ty'))
        in
        let arg, ty_arg = elaborate' arg env in
        let arg_asc = ascribe arg (interior ty_arg ty) ty in
        let elaborate_pattern (pattern : Typed_ast.pattern) =
          let case : Type.case =
            List.find
              (fun (Type.Case (name, _)) -> name = pattern.casename)
              cases
          in
          let (Type.Case (casename, tys)) = case in
          if List.length pattern.args <> List.length tys then
            raise (err "pattern does not match variant type");
          let env' =
            List.fold_left2
              (fun env arg ty -> Env.add arg (interior ty ty, ty) env)
              env pattern.args tys
          in
          let body, ty_body = elaborate' pattern.body env' in
          (Ast.{ casename; args = pattern.args; body }, ty_body)
        in
        let patterns, tys = List.split (List.map elaborate_pattern patterns) in
        let ty_meet = List.fold_left Ev.meet Type.Any tys in
        let ev = interior ty_meet ty_meet in
        (make (Match (arg_asc, patterns, ty)) ev, ty_meet)
    | Seq args ->
        let args, tys =
          List.split (List.map (fun arg -> elaborate' arg env) args)
        in
        let ty = List.hd (List.rev tys) in
        let ev = interior ty ty in
        (make (Seq args) ev, ty)

  and ascribe (t : Ev.t Ast.t) (ev : Ev.t) (ty : Type.t) : Ev.t Ast.t =
    Ast.{ desc = Ascribe (t, ty); meta = ev }
end

module TypeElaborate = Elaborate (TypeEvidence)
module RuntimeTypeElaborate = Elaborate (RuntimeTypeEvidence)
