(*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

open Type
open TypeUtil
open Constraint

let merge_trust_var constr =
  Trust_constraint.(
    match constr with
    | TrustResolved t -> t
    | TrustUnresolved bound -> get_trust bound |> Trust.fix
  )

(****************** signature contexts *********************)

(* Once a context is merged with other contexts it depends on, it is ready to be
   optimized for use by contexts that in turn depend on it. The only interesting
   part of the context for such use is its exports. Thus, we call a context
   optimized for such use a "signature context."

   A signature context should contain only descriptions of its exports. Anything
   that does not contribute to those descriptions are redundant and should be
   discarded: otherwise, we would end up paying significant hidden costs, in
   terms of memory (storing in the heap) and time (serializing / deserializing
   from the heap, garbage collection pauses).

   As it turns out, it is not always possible or even desirable to describe the
   exports in "closed form" as types: instead, the descriptions take the form of
   a (mutually) recursive system of "equations" relating type variables to
   "solved" types, and the types themselves may contain these type variables
   inside. As indicated above, such descriptions are both necessary (due to
   recursive types) and desirable (due to sharing, and for cyclically depending
   contexts, due to the need to deal with the exports of all such contexts at
   once). This form is quite adequate for type checking (indeed, it can be
   thought of as merely a condensed form of usual contexts, which already
   exhibit this "interlinked" behavior).

   Unsurprisingly, since types also make references to some other maps in a
   context (property maps, envs), these references need to be carried around
   as well for the descriptions to be complete.

   We can collect compact and complete descriptions of exports by using a custom
   type visitor that marks relevant things that need to be carried over as it
   walks the context from its exports. Once the walk is done, we can replace the
   corresponding parts of the context (graph, property maps, envs) by their
   reduced forms.

   NOTE: Changing ids of entities stored in the Context is dangerous here,
   since it would require changing every occurence of that id. As such, the mapper
   currently only traverses ids for side-effects, but returns the original id.

   There are also several places where we add something to a map, recurse, then
   readd. The original add makes sure we don't visit the same entity twice, the
   second actually changes the entity in the exports.
*)

class context_optimizer ~no_lowers =
  object (self)
    inherit [Polarity.t] Type_mapper.t as super

    val mutable reduced_graph : Type.Constraint.graph = IMap.empty

    val mutable reduced_trust_graph = IMap.empty

    val mutable reduced_property_maps = Properties.Map.empty

    val mutable reduced_call_props = IMap.empty

    val mutable reduced_export_maps = Exports.Map.empty

    val mutable reduced_evaluated = Eval.Map.empty

    val mutable export_reason = None

    method tvar cx pole r id =
      let (root_id, _) = Context.find_constraints cx id in
      if id == root_id then (
        if IMap.mem id reduced_graph then
          id
        else
          let t = lazy (self#type_ cx pole (Flow_js_utils.merge_tvar ~no_lowers cx r id)) in
          let node = create_root (FullyResolved t) in
          reduced_graph <- IMap.add id node reduced_graph;
          ignore (Lazy.force t);
          id
      ) else (
        ignore (self#tvar cx pole r root_id);
        let node = create_goto root_id in
        reduced_graph <- IMap.add id node reduced_graph;
        id
      )

    method trust_var cx pole id =
      let (root_id, constr) = Context.find_trust_constraints cx id in
      if id == root_id then (
        if IMap.mem id reduced_trust_graph then
          id
        else
          let t = merge_trust_var constr in
          let node = Trust_constraint.new_resolved_root t in
          reduced_trust_graph <- IMap.add id node reduced_trust_graph;
          id
      ) else (
        ignore (self#trust_var cx pole root_id);
        let node = Trust_constraint.TrustGoto root_id in
        reduced_trust_graph <- IMap.add id node reduced_trust_graph;
        id
      )

    method props cx pole id =
      if Properties.Map.mem id reduced_property_maps then
        id
      else
        let pmap = Context.find_props cx id in
        reduced_property_maps <- Properties.Map.add id pmap reduced_property_maps;
        let pmap' = NameUtils.Map.ident_map (self#prop cx pole) pmap in
        reduced_property_maps <- Properties.Map.add id pmap' reduced_property_maps;
        id

    method call_prop cx pole id =
      if IMap.mem id reduced_call_props then
        id
      else
        let t = Context.find_call cx id in
        reduced_call_props <- IMap.add id t reduced_call_props;
        let t' = self#type_ cx pole t in
        reduced_call_props <- IMap.add id t' reduced_call_props;
        id

    method! export_types cx map_cx e =
      let { exports_tmap; cjs_export; has_every_named_export } = e in
      let exports_tmap' = self#exports cx map_cx exports_tmap in
      let cjs_export' =
        OptionUtils.ident_map
          (fun exp ->
            export_reason <- Some (reason_of_t exp);
            self#type_ cx map_cx exp)
          cjs_export
      in
      if exports_tmap == exports_tmap' && cjs_export == cjs_export' then
        e
      else
        { exports_tmap = exports_tmap'; cjs_export = cjs_export'; has_every_named_export }

    method exports cx pole id =
      if Exports.Map.mem id reduced_export_maps then
        id
      else
        let tmap = Context.find_exports cx id in
        let map_pair p =
          let (loc, t) = p in
          export_reason <- Some (reason_of_t t);
          let t' = self#type_ cx pole t in
          if t == t' then
            p
          else
            (loc, t')
        in
        reduced_export_maps <- Exports.Map.add id tmap reduced_export_maps;
        let tmap' = NameUtils.Map.ident_map map_pair tmap in
        reduced_export_maps <- Exports.Map.add id tmap' reduced_export_maps;
        id

    method eval_id cx pole id =
      if Eval.Map.mem id reduced_evaluated then
        id
      else
        match Eval.Map.find_opt id (Context.evaluated cx) with
        | None -> id
        | Some t ->
          reduced_evaluated <- Eval.Map.add id t reduced_evaluated;
          let t' = self#type_ cx pole t in
          reduced_evaluated <- Eval.Map.add id t' reduced_evaluated;
          id

    method! destructor cx map_cx t =
      match t with
      | ReactCheckComponentConfig map ->
        let map' = NameUtils.Map.ident_map (self#prop cx map_cx) map in
        if map' == map then
          t
        else
          ReactCheckComponentConfig map'
      | NonMaybeType -> t
      | ReactCheckComponentRef -> t
      | PropertyType _ -> t
      | ElementType { index_type } ->
        let index_type' = self#type_ cx map_cx index_type in
        if index_type' == index_type then
          t
        else
          ElementType { index_type = index_type' }
      | OptionalIndexedAccessNonMaybeType { index = OptionalIndexedAccessTypeIndex index_type } ->
        let index_type' = self#type_ cx map_cx index_type in
        if index_type' == index_type then
          t
        else
          OptionalIndexedAccessNonMaybeType { index = OptionalIndexedAccessTypeIndex index_type' }
      | OptionalIndexedAccessNonMaybeType { index = OptionalIndexedAccessStrLitIndex _ } -> t
      | OptionalIndexedAccessResultType _ -> t
      | ReadOnlyType -> t
      | RequiredType
      | PartialType ->
        t
      | SpreadType (options, tlist, acc) ->
        let tlist' = ListUtils.ident_map (self#object_kit_spread_operand cx map_cx) tlist in
        let acc' = OptionUtils.ident_map (self#object_kit_spread_operand_slice cx map_cx) acc in
        if tlist' == tlist && acc == acc' then
          t
        else
          SpreadType (options, tlist', acc')
      | RestType (options, x) ->
        let x' = self#type_ cx map_cx x in
        if x' == x then
          t
        else
          RestType (options, x')
      | ValuesType -> t
      | CallType { from_maptype; args } ->
        let args' = ListUtils.ident_map (self#type_ cx map_cx) args in
        if args' == args then
          t
        else
          CallType { from_maptype; args = args' }
      | ConditionalType { distributive_tparam_name; infer_tparams; extends_t; true_t; false_t } ->
        let infer_tparams' = ListUtils.ident_map (self#type_param cx map_cx) infer_tparams in
        let extends_t' = self#type_ cx map_cx extends_t in
        let true_t' = self#type_ cx map_cx true_t in
        let false_t' = self#type_ cx map_cx false_t in
        if
          infer_tparams' == infer_tparams
          && extends_t' == extends_t
          && true_t' == true_t
          && false_t' == false_t
        then
          t
        else
          ConditionalType
            {
              distributive_tparam_name;
              infer_tparams = infer_tparams';
              extends_t = extends_t';
              true_t = true_t';
              false_t = false_t';
            }
      | TypeMap tmap ->
        let tmap' = self#type_map cx map_cx tmap in
        if tmap' == tmap then
          t
        else
          TypeMap tmap'
      | ReactConfigType default_props ->
        let default_props' = self#type_ cx map_cx default_props in
        if default_props' == default_props then
          t
        else
          ReactConfigType default_props'
      | MappedType { property_type; mapped_type_flags; homomorphic; distributive_tparam_name } ->
        let property_type' = self#type_ cx map_cx property_type in
        if property_type' == property_type then
          t
        else
          MappedType
            {
              property_type = property_type';
              mapped_type_flags;
              homomorphic;
              distributive_tparam_name;
            }
      | ReactElementPropsType
      | ReactElementConfigType
      | ReactElementRefType
      | IdxUnwrapType ->
        t
      | LatentPred (t_in, i) ->
        let t_in' = self#type_ cx map_cx t_in in
        if t_in' == t_in then
          t
        else
          LatentPred (t_in', i)

    method! type_ cx pole t =
      begin
        match t with
        | DefT (_, trust, _) when Context.trust_tracking cx && is_ident trust ->
          ignore (self#trust_var cx pole (as_ident trust))
        | _ -> ()
      end;
      match t with
      | InternalT _ -> Utils_js.assert_false "internal types should not appear in signatures"
      | UnionT (_, rep) ->
        if not (UnionRep.is_optimized_finally rep) then
          UnionRep.optimize
            rep
            ~reasonless_eq:TypeUtil.reasonless_eq
            ~flatten:(Type_mapper.union_flatten cx)
            ~find_resolved:(Context.find_resolved cx)
            ~find_props:(Context.find_props cx);
        super#type_ cx pole t
      | _ -> super#type_ cx pole t

    method get_reduced_graph = reduced_graph

    method get_reduced_trust_graph = reduced_trust_graph

    method get_reduced_property_maps = reduced_property_maps

    method get_reduced_call_props = reduced_call_props

    method get_reduced_export_maps = reduced_export_maps

    method get_reduced_evaluated = reduced_evaluated
  end
