(* Removes unnecessary ascriptions. When the statically known evidence is more precise than the evidence being ascribed we can remove that evidence. *)
(*
   If the ascription could result in a runtime error, we can not remove it.
   Thus, we can only remove the ascriptions that we statically know won't have any effect.contents
   This will ocurr when we ascribe to the same evidence, or a less precise evidence.

   Ascribing to ? can always be removed, as it does nothing.

   Example 1:

   let x = 10 :: int in
   x :: int * 2

   In this example the first cast can be removed since integers will always have int evidence.
   The second cast can also be removed, because we know statically that x has integer evidence.

   Example 2:

   (fun x : int -> x + x) 10
   elaborates to
   ((fun x -> (x :: int) + (x :: int)) :: int -> int)(10 :: int)
*)

open Evidence

module Make (Ev : Evidence with type g = Type.t) = struct
  let can_remove_annotation (value_ev : Ev.t) (ascription_ev : Ev.t) =
    (* an ascription where we statically know exactly the type of the
    evidence being used to ascribe can be removed when we know that the
    term is more precise, and thus it will never result in a runtime
      error. *)
    let ev_opt = Ev.consistent_transitivity value_ev ascription_ev in

    match ev_opt with
    | Some ev ->
        (*Format.printf "Checking removal of annotation: value_ev=%a, ascription_ev=%a, meet=%a\n"
      Ev.pp value_ev Ev.pp ascription_ev Ev.pp ev;*)
        value_ev = ev
    | None -> false

  let rec simplify env (t : Ev.t Ast.t) =
    let open Ast in
    let mk meta desc = { desc; meta } in
    match t.desc with
    | Const _ -> t
    | Var _ -> t
    | Lam (args, body) -> mk t.meta @@ Lam (args, simplify env body)
    | App (t_f, t_args) ->
        mk t.meta @@ App (simplify env t_f, List.map (simplify env) t_args)
    | Prim1 (op, t1) -> mk t.meta @@ Prim1 (op, simplify env t1)
    | Prim2 (op, t1, t2) ->
        mk t.meta @@ Prim2 (op, simplify env t1, simplify env t2)
    | If (t1, t2, t3) ->
        mk t.meta @@ If (simplify env t1, simplify env t2, simplify env t3)
    | Let (binding, body) ->
        mk t.meta
        @@ Let
             ( { binding with term = simplify env binding.term },
               simplify env body )
    | Letrec (bindings, body) ->
        let bindings =
          List.map (fun b -> { b with term = simplify env b.term }) bindings
        in
        mk t.meta @@ Letrec (bindings, simplify env body)
    | Tuple terms -> mk t.meta @@ Tuple (List.map (simplify env) terms)
    | Get (t1, i) -> mk t.meta @@ Get (simplify env t1, i)
    | Vect (t1, t2) -> mk t.meta @@ Vect (simplify env t1, simplify env t2)
    | VectGet (t1, t2) -> mk t.meta @@ VectGet (simplify env t1, simplify env t2)
    | VectSet (t1, t2, t3) ->
        mk t.meta @@ VectSet (simplify env t1, simplify env t2, simplify env t3)
    | VectLength t1 -> mk t.meta @@ VectLength (simplify env t1)
    | For (x, t1, t2, body) ->
        mk t.meta @@ For (x, simplify env t1, simplify env t2, simplify env body)
    | Variant (name, terms, ty) ->
        mk t.meta @@ Variant (name, List.map (simplify env) terms, ty)
    | Match (term, patterns, ty) ->
        let patterns =
          List.map (fun p -> { p with body = simplify env p.body }) patterns
        in
        mk t.meta @@ Match (simplify env term, patterns, ty)
    | Seq terms -> mk t.meta @@ Seq (List.map (simplify env) terms)
    | Ascribe (t1, ty) ->
        let t1 = simplify env t1 in
        if can_remove_annotation t1.meta t.meta then t1
        else mk t.meta @@ Ascribe (simplify env t1, ty)
end

module SimplifyTypeEvidence = Make (TypeEvidence)
module SimplifyRuntimeTypeEvidence = Make (RuntimeTypeEvidence)
