module ElaborateDyn
    (Ev : Evidence.Evidence with type g = Type.t)
    (TM : Evidence.TypeMatching with type t = Ev.t) =
struct
  open Anf
  open Common

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

  let ascribe_dyn (c : Ev.t cexpr) (a : Ev.t atom) (inv : Inversion.t) =
    let ev_value = c.cexpr_meta in
    let ev_ascription = a.atom_meta in
    let ev_inverted =
      match inv with
      | Inversion.IDom i -> TM.dom ev_ascription i
      | Inversion.IProd i -> TM.proj ev_ascription i
      | Inversion.ICod -> TM.cod ev_ascription
      | Inversion.IVec variance -> (
          match variance with
          | Covariant -> TM.vec ev_ascription
          | Contravariant -> Option.map Ev.invert (TM.vec ev_ascription))
    in
    match ev_inverted with
    | Some ev
      when Ev.can_be_optimized ev_value
           && Ev.can_be_optimized ev
           && Ev.consistent_transitivity ev_value ev = Option.some ev_value ->
        c
    | Some ev -> cexpr_of_desc (AscribeDyn (c, a, inv)) ev
    | None -> assert false

  let rec elaborate (expr : Ev.t expr) : Ev.t expr =
    match expr.expr_desc with
    | Ret cexpr ->
        let* c = elaborate_cexpr cexpr in
        ret c expr.expr_meta
    | Let (binding, body) ->
        let* bindings = elaborate_bindings [ binding ] in
        bind (List.hd bindings) (elaborate body)
    | Letrec (bindings, body) ->
        let* bindings' = elaborate_bindings bindings in
        bind_rec bindings' (elaborate body)
    | Seq exprs -> expr_of_desc (Seq (List.map elaborate exprs)) expr.expr_meta

  and ascribe_ret (expr : Ev.t expr) ev_lambda =
    match expr.expr_desc with
    | Ret c ->
        let* c = elaborate_cexpr c in
        let ascribed' =
          ascribe_dyn c (atom_of_desc (Var "self") ev_lambda) Inversion.ICod
        in
        expr_of_desc (Ret ascribed') expr.expr_meta
    | Let (binding, e) ->
        let* bindings = elaborate_bindings [ binding ] in
        bind (List.hd bindings) (ascribe_ret e ev_lambda)
    | Letrec (bindings, e) ->
        let* bindings' = elaborate_bindings bindings in
        bind_rec bindings' (ascribe_ret e ev_lambda)
    | Seq exprs ->
        let rec go acc exprs =
          match exprs with
          | [] -> List.rev acc
          | [ expr ] -> go (ascribe_ret expr ev_lambda :: acc) []
          | expr :: tl -> go (elaborate expr :: acc) tl
        in
        expr_of_desc (Seq (go [] exprs)) expr.expr_meta

  and elaborate_cexpr (cexpr : Ev.t cexpr) (k : Ev.t cexpr -> Ev.t expr) :
      Ev.t expr =
    let meta = cexpr.cexpr_meta in
    match cexpr.cexpr_desc with
    | Atom a -> k @@ cexpr_of_desc (Atom (elaborate_atom a)) meta
    | App (f, args) ->
        let f' = elaborate_atom f in
        let* args' = elaborate_args f' args 0 in
        k @@ cexpr_of_desc (App (f', args')) meta
    | Lam (args, body) ->
        k @@ cexpr_of_desc (Lam (args, ascribe_ret body meta)) meta
    | Prim1 (op, c) ->
        let* c' = elaborate_cexpr c in
        k @@ cexpr_of_desc (Prim1 (op, c')) meta
    | Prim2 (op, a1, a2) ->
        k
        @@ cexpr_of_desc (Prim2 (op, elaborate_atom a1, elaborate_atom a2)) meta
    | If (cond, then_branch, else_branch) ->
        k
        @@ cexpr_of_desc
             (If
                ( elaborate_atom cond,
                  elaborate then_branch,
                  elaborate else_branch ))
             meta
    | Match (scrutinee, patterns, ty) ->
        k
        @@ cexpr_of_desc
             (Match
                ( elaborate_atom scrutinee,
                  List.map elaborate_pattern patterns,
                  ty ))
             meta
    | Tuple atoms ->
        k @@ cexpr_of_desc (Tuple (List.map elaborate_atom atoms)) meta
    | Get (c, idx) ->
        let* c' = elaborate_cexpr c in
        let tuple_name = Gensym.fresh "tuple" in
        let tuple_binding = { name = tuple_name; cexpr = c' } in
        let tuple_atom = atom_of_desc (Var tuple_name) c'.cexpr_meta in
        let tuple_cexpr = cexpr_of_desc (Atom tuple_atom) c'.cexpr_meta in
        let get_result = cexpr_of_desc (Get (tuple_cexpr, idx)) meta in
        let ascribed =
          ascribe_dyn get_result tuple_atom (Inversion.IProd idx)
        in
        bind tuple_binding (k ascribed)
    | Vect (a1, a2) ->
        k @@ cexpr_of_desc (Vect (elaborate_atom a1, elaborate_atom a2)) meta
    | VectGet (a_vec, a_idx) ->
        let a_vec' = elaborate_atom a_vec in
        let a_idx' = elaborate_atom a_idx in
        k
        @@ ascribe_dyn
             (cexpr_of_desc (VectGet (a_vec', a_idx')) meta)
             a_vec'
             Inversion.(IVec Covariant)
    | VectSet (a_vec, a_idx, a_val) ->
        let a_vec' = elaborate_atom a_vec in
        let a_idx' = elaborate_atom a_idx in
        let a_val' = elaborate_atom a_val in
        let ascribed_val =
          ascribe_dyn
            (cexpr_of_desc (Atom a_val') a_val'.atom_meta)
            a_vec'
            Inversion.(IVec Contravariant)
        in
        let name = Gensym.fresh "val_asc" in
        let binding = { name; cexpr = ascribed_val } in
        let ascribed_val = atom_of_desc (Var name) meta in
        bind binding
          (k @@ cexpr_of_desc (VectSet (a_vec', a_idx', ascribed_val)) meta)
    | VectLength c ->
        let* c' = elaborate_cexpr c in
        k @@ cexpr_of_desc (VectLength c') meta
    | Ascribe (c, ty) ->
        let* c' = elaborate_cexpr c in
        k @@ cexpr_of_desc (Ascribe (c', ty)) meta
    | AscribeDyn (c, a, inv) ->
        let* c' = elaborate_cexpr c in
        k @@ ascribe_dyn c' a inv
    | For (var, start, stop, body) ->
        k
        @@ cexpr_of_desc
             (For
                (var, elaborate_atom start, elaborate_atom stop, elaborate body))
             meta
    | Variant (name, atoms, ty) ->
        k
        @@ cexpr_of_desc
             (Variant (name, List.map elaborate_atom atoms, ty))
             meta

  and elaborate_arg (f : Ev.t atom) (idx : int) (arg : Ev.t atom) (meta : Ev.t)
      (k : Ev.t atom -> Ev.t expr) : Ev.t expr =
    let arg' = elaborate_atom arg in
    let ascribed =
      ascribe_dyn (cexpr_of_desc (Atom arg') meta) f (Inversion.IDom idx)
    in
    let name = Gensym.fresh (Format.sprintf "arg_%d" idx) in
    let binding = { name; cexpr = ascribed } in
    let atom = atom_of_desc (Var name) meta in
    bind binding (k atom)

  and elaborate_args (f : Ev.t atom) (args : Ev.t atom list) (idx : int)
      (k : Ev.t atom list -> Ev.t expr) : Ev.t expr =
    let rec go acc idx args =
      match args with
      | [] -> k (List.rev acc)
      | arg :: rest ->
          let* arg_atom = elaborate_arg f idx arg arg.atom_meta in
          go (arg_atom :: acc) (idx + 1) rest
    in
    go [] idx args

  and elaborate_atom (atom : Ev.t atom) : Ev.t atom =
    match atom.atom_desc with Const _ | Var _ -> atom

  and elaborate_pattern (pattern : Ev.t pattern) : Ev.t pattern =
    { pattern with body = elaborate pattern.body }

  and elaborate_bindings (bindings : Ev.t binding list)
      (k : Ev.t binding list -> Ev.t expr) : Ev.t expr =
    let rec go acc bindings =
      match bindings with
      | [] -> k (List.rev acc)
      | binding :: tl ->
          let* cexpr' = elaborate_cexpr binding.cexpr in
          let binding' = { binding with cexpr = cexpr' } in
          go (binding' :: acc) tl
    in
    go [] bindings
end

module TypeElaborateDyn =
  ElaborateDyn (Evidence.TypeEvidence) (Evidence.TypeMatching)

module RuntimeTypeElaborateDyn =
  ElaborateDyn (Evidence.RuntimeTypeEvidence) (Evidence.RuntimeTypeMatching)
