(*
 * 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.
 *)

module Ast = Flow_ast
module Flow = Flow_js
open Reason
open Type
open TypeUtil
include Func_sig_intf

class func_scope_visitor
  cx ~has_return_annot ~ret_annot_loc ~return_t ~yield_t ~next_t ~body_loc kind exhaust =
  object (this)
    inherit
      [ALoc.t, ALoc.t * Type.t, ALoc.t, ALoc.t * Type.t] Flow_polymorphic_ast_mapper.mapper as super

    val mutable no_return = true

    val mutable has_throw = false

    val mutable no_yield = true

    method on_type_annot x = x

    method on_loc_annot x = x

    method! function_ fn = fn

    method! component_declaration c = c

    method visit statements =
      ignore @@ this#statement_list statements;
      if not has_return_annot then (
        ( if no_return && has_throw then
          let reason = mk_reason REmpty body_loc in
          let t = EmptyT.make reason in
          Flow.flow cx (t, UseT (unknown_use, return_t))
        );
        if no_yield then Flow.flow_t cx (VoidT.make (mk_reason RVoid body_loc), yield_t)
      )

    method! switch ({ Ast.Statement.Switch.exhaustive_out = (loc, t); _ } as switch) =
      Base.Option.iter exhaust ~f:(fun (exhaustive_ts, exhaust_locs, _) ->
          if Base.List.mem ~equal:ALoc.equal exhaust_locs loc then
            exhaustive_ts := t :: !exhaustive_ts
      );
      super#switch switch

    method! yield ({ Ast.Expression.Yield.result_out = (_, t); argument; delegate; _ } as yield) =
      let use_op =
        if delegate then
          unknown_use
        else
          Op
            (GeneratorYield
               {
                 value =
                   Base.Option.value_map argument ~default:(reason_of_t t) ~f:(fun expr ->
                       mk_expression_reason (Typed_ast_utils.untyped_ast_mapper#expression expr)
                   );
               }
            )
      in
      Flow.flow cx (t, UseT (use_op, yield_t));
      no_yield <- false;
      super#yield yield

    (* Override statement so that we have the loc for return *)
    method! statement (loc, stmt) =
      begin
        match stmt with
        | Ast.Statement.Return return -> this#custom_return loc return
        | Ast.Statement.Throw _ -> has_throw <- true
        | _ -> ()
      end;
      super#statement (loc, stmt)

    method! call loc expr =
      let open Ast.Expression.Call in
      let { callee; arguments; _ } = expr in
      ( if Flow_ast_utils.is_call_to_invariant callee then
        match arguments with
        (* invariant() and invariant(false, ...) are treated like throw *)
        | (_, { Ast.Expression.ArgList.arguments = []; comments = _ })
        | ( _,
            {
              Ast.Expression.ArgList.arguments =
                Ast.Expression.Expression
                  (_, Ast.Expression.BooleanLiteral { Ast.BooleanLiteral.value = false; _ })
                :: _;
              comments = _;
            }
          ) ->
          has_throw <- true
        | _ -> ()
      );
      super#call loc expr

    method custom_return loc { Ast.Statement.Return.return_out = (_, t); argument; _ } =
      let open Func_class_sig_types.Func in
      let t =
        match kind with
        | Async ->
          (* Convert the return expression's type T to Promise<T>. If the
             * expression type is itself a Promise<T>, ensure we still return
             * a Promise<T> via Promise.resolve. *)
          let async_return_reason = mk_reason RAsyncReturn loc in
          let t' =
            Flow.get_builtin_typeapp
              cx
              (mk_reason (desc_of_t t) loc)
              "Promise"
              [Type_operation_utils.Promise.await cx async_return_reason t]
          in
          Flow.reposition cx loc t'
        | Generator _ ->
          (* Convert the return expression's type R to Generator<Y,R,N>, where
           * Y and R are internals, installed earlier. *)
          let t' =
            Flow.get_builtin_typeapp
              cx
              (mk_reason (desc_of_t t) loc)
              "Generator"
              [
                yield_t;
                Tvar.mk_where cx (mk_reason (RCustom "generator return") loc) (fun tvar ->
                    Flow.flow_t cx (t, tvar)
                );
                next_t;
              ]
          in
          Flow.reposition cx loc t'
        | AsyncGenerator _ ->
          let t' =
            Flow.get_builtin_typeapp
              cx
              (mk_reason (desc_of_t t) loc)
              "AsyncGenerator"
              [
                yield_t;
                Tvar.mk_where cx (mk_reason (RCustom "async generator return") loc) (fun tvar ->
                    Flow.flow_t cx (t, tvar)
                );
                next_t;
              ]
          in
          Flow.reposition cx loc t'
        | _ -> t
      in
      let use_op =
        Op
          (FunReturnStatement
             {
               value =
                 Base.Option.value_map argument ~default:(reason_of_t t) ~f:(fun expr ->
                     mk_expression_reason (Typed_ast_utils.untyped_ast_mapper#expression expr)
                 );
             }
          )
      in
      if not has_return_annot then Context.add_missing_local_annot_lower_bound cx ret_annot_loc t;
      Flow.flow cx (t, UseT (use_op, return_t));
      no_return <- false
  end

module Make
    (Statement : Statement_sig.S)
    (CT : Func_class_sig_types.Config.S)
    (C : Func_params_intf.Config with module Types := CT)
    (F : Func_params_intf.S with module Config_types := CT and module Config := C)
    (T : Func_class_sig_types.Func.S with module Config := CT and module Param := F.Types) :
  S with module Config_types := CT and module Config := C and module Param := F and module Types = T =
struct
  module Types = T
  open Func_class_sig_types

  let this_param = F.this

  let default_constructor reason =
    {
      T.reason;
      kind = Func.Ctor;
      tparams = None;
      fparams = F.empty (fun _ _ _ -> None);
      body = None;
      return_t = Annotated (VoidT.why reason);
      effect_ = ArbitraryEffect;
      ret_annot_loc = Reason.loc_of_reason reason;
      statics = None;
    }

  let field_initializer reason expr annot_loc return_annot_or_inferred =
    {
      T.reason;
      kind = Func.FieldInit expr;
      tparams = None;
      fparams = F.empty (fun _ _ _ -> None);
      body = None;
      return_t = return_annot_or_inferred;
      effect_ = ArbitraryEffect;
      ret_annot_loc = annot_loc;
      statics = None;
    }

  let functiontype cx ~arrow func_loc this_default x =
    let { T.reason; kind; tparams; fparams; return_t; statics; effect_; _ } = x in
    let this_type = F.this fparams |> Base.Option.value ~default:this_default in
    let return_t =
      match return_t with
      | Inferred t
        when Context.typing_mode cx <> Context.CheckingMode
             && (not @@ F.all_params_annotated x.T.fparams) ->
        Context.mk_placeholder cx (TypeUtil.reason_of_t t)
      | _ -> TypeUtil.type_t_of_annotated_or_inferred return_t
    in
    let type_guard =
      match kind with
      | Func.TypeGuard p -> Some p
      | _ -> None
    in
    let funtype =
      {
        Type.this_t = (this_type, This_Function);
        params = F.value fparams;
        rest_param = F.rest fparams;
        return_t;
        effect_;
        type_guard;
        def_reason = reason;
      }
    in
    let statics_t =
      match statics with
      | Some t -> t
      | None -> Obj_type.mk_with_proto cx reason (Type.FunProtoT reason) ~obj_kind:Type.Inexact
    in
    let t = DefT (reason, FunT (statics_t, funtype)) in
    if not arrow then Base.Option.iter func_loc ~f:(Type_env.bind_function_this cx this_type);
    poly_type_of_tparams (Type.Poly.generate_id ()) tparams t

  let methodtype
      cx method_this_loc this_default { T.reason; kind; tparams; fparams; return_t; effect_; _ } =
    let params = F.value fparams in
    let (params_names, params_tlist) = List.split params in
    let rest_param = F.rest fparams in
    let def_reason = reason in
    let this_anno_t = F.this fparams in
    Base.Option.both this_anno_t method_this_loc
    |> Base.Option.iter ~f:(fun (t, loc) -> Type_env.bind_function_this cx t loc);
    let param_this_t = Base.Option.value ~default:this_default this_anno_t in
    let type_guard =
      match kind with
      | Func.TypeGuard g -> Some g
      | _ -> None
    in
    let t =
      DefT
        ( reason,
          FunT
            ( dummy_static reason,
              mk_boundfunctiontype
                ~this:param_this_t
                ~effect_
                params_tlist
                ~rest_param
                ~def_reason
                ~params_names
                ~type_guard
                (TypeUtil.type_t_of_annotated_or_inferred return_t)
            )
        )
    in
    poly_type_of_tparams (Type.Poly.generate_id ()) tparams t

  let gettertype ({ T.return_t; _ } : T.t) = TypeUtil.type_t_of_annotated_or_inferred return_t

  let settertype { T.reason; fparams; _ } =
    match F.value fparams with
    | (_, param_t) :: _ ->
      (* setters must have exactly one param. more than one is a syntax error,
         so ignore any extras *)
      param_t
    | [] -> AnyT.error reason

  let toplevels cx x =
    let { T.reason = reason_fn; kind; fparams; body; ret_annot_loc; return_t; _ } = x in
    let body_loc =
      let open Ast.Function in
      match body with
      | Some (BodyBlock (loc, _)) -> loc
      | Some (BodyExpression (loc, _)) -> loc
      | None -> ALoc.none
    in
    let reason = mk_reason RFunctionBody body_loc in

    (* Set the scope early so default exprs can reference earlier params *)
    let prev_scope_kind =
      let var_scope_kind =
        match kind with
        | Func.Ordinary
        | Func.FieldInit _
        | Func.TypeGuard _ ->
          Name_def.Ordinary
        | Func.Async -> Name_def.Async
        | Func.Generator _ -> Name_def.Generator
        | Func.AsyncGenerator _ -> Name_def.AsyncGenerator
        | Func.Ctor -> Name_def.Ctor
      in
      Type_env.set_scope_kind cx var_scope_kind
    in

    (* add param bindings *)
    let params_ast = F.eval cx fparams in

    let (yield_t, next_t) =
      match kind with
      | Func.Generator _
      | Func.AsyncGenerator _ ->
        let yield_t = Tvar.mk cx (replace_desc_reason (RCustom "yield") reason) in
        let next_t =
          match return_t with
          | Annotated _ -> Tvar.mk cx (replace_desc_reason RNext reason)
          | Inferred _ -> VoidT.make (replace_desc_reason RUnannotatedNext reason)
        in
        let return_targ = Tvar.mk cx reason in
        let (iterable, generator) =
          match kind with
          | Func.Generator _ -> ("$Iterable", "Generator")
          | Func.AsyncGenerator _ -> ("$AsyncIterable", "AsyncGenerator")
          | _ -> failwith "Bad kind"
        in
        let () =
          let t = Flow.get_builtin_typeapp cx reason iterable [yield_t; return_targ; next_t] in
          let t =
            Flow.reposition
              cx
              (type_t_of_annotated_or_inferred return_t |> reason_of_t |> loc_of_reason)
              t
          in
          Flow.flow_t cx (type_t_of_annotated_or_inferred return_t, t)
        in
        let () =
          let t = Flow.get_builtin_typeapp cx reason generator [yield_t; return_targ; next_t] in
          let t =
            Flow.reposition
              cx
              (type_t_of_annotated_or_inferred return_t |> reason_of_t |> loc_of_reason)
              t
          in
          Flow.flow_t cx (t, type_t_of_annotated_or_inferred return_t)
        in
        (yield_t, next_t)
      | _ ->
        ( DefT (replace_desc_reason (RCustom "no yield") reason, MixedT Mixed_everything),
          DefT (replace_desc_reason (RCustom "no next") reason, MixedT Mixed_everything)
        )
    in

    let (statements, reconstruct_body, maybe_implicit_void_return) =
      let open Ast.Statement in
      match body with
      | None -> ([], Fun.const None, true)
      | Some (Ast.Function.BodyBlock (loc, { Block.body; comments })) ->
        (body, (fun body -> Some (Ast.Function.BodyBlock (loc, { Block.body; comments }))), true)
      | Some (Ast.Function.BodyExpression expr) ->
        ( [
            ( fst expr,
              Return { Return.argument = Some expr; comments = None; return_out = fst expr }
            );
          ],
          (function
          | [(_, Return { Return.argument = Some expr; comments = _; return_out = _ })]
          | [(_, Expression { Expression.expression = expr; _ })] ->
            Some (Ast.Function.BodyExpression expr)
          | _ -> failwith "expected return body"),
          false
        )
    in
    let (has_return_annot, return_t) =
      match return_t with
      | Inferred t -> (false, t)
      | Annotated t -> (true, t)
    in

    (* statement visit pass *)
    let statements_ast = Statement.statement_list cx statements in

    let body_ast = reconstruct_body statements_ast in
    (* build return type for void funcs *)
    let (init_ast, exhaust) =
      if maybe_implicit_void_return then
        let loc = loc_of_t return_t in
        (* Some branches add an ImplicitTypeParam frame to force our flow_use_op
         * algorithm to pick use_ops outside the provided loc. *)
        let (use_op, void_t, init_ast) =
          match kind with
          | Func.Ordinary
          | Func.Ctor ->
            let t = VoidT.at loc in
            let use_op =
              Op
                (FunImplicitReturn
                   { fn = reason_fn; upper = reason_of_t return_t; type_guard = false }
                )
            in
            (use_op, t, None)
          | Func.Async ->
            let reason = mk_annot_reason (RType (OrdinaryName "Promise")) loc in
            let void_t = VoidT.at loc in
            let t = Flow.get_builtin_typeapp cx reason "Promise" [void_t] in
            let use_op =
              Op
                (FunImplicitReturn
                   { fn = reason_fn; upper = reason_of_t return_t; type_guard = false }
                )
            in
            let use_op = Frame (ImplicitTypeParam, use_op) in
            (use_op, t, None)
          | Func.Generator _ ->
            let reason = mk_annot_reason (RType (OrdinaryName "Generator")) loc in
            let void_t = VoidT.at loc in
            let t = Flow.get_builtin_typeapp cx reason "Generator" [yield_t; void_t; next_t] in
            let use_op =
              Op
                (FunImplicitReturn
                   { fn = reason_fn; upper = reason_of_t return_t; type_guard = false }
                )
            in
            let use_op = Frame (ImplicitTypeParam, use_op) in
            (use_op, t, None)
          | Func.AsyncGenerator _ ->
            let reason = mk_annot_reason (RType (OrdinaryName "AsyncGenerator")) loc in
            let void_t = VoidT.at loc in
            let t = Flow.get_builtin_typeapp cx reason "AsyncGenerator" [yield_t; void_t; next_t] in
            let use_op =
              Op
                (FunImplicitReturn
                   { fn = reason_fn; upper = reason_of_t return_t; type_guard = false }
                )
            in
            let use_op = Frame (ImplicitTypeParam, use_op) in
            (use_op, t, None)
          | Func.FieldInit e ->
            let (((_, t), _) as ast) =
              Statement.expression ~encl_ctx:Enclosing_context.NoContext cx e
            in
            let body = mk_expression_reason e in
            let use_op = Op (InitField { op = reason_fn; body }) in
            (use_op, t, Some ast)
          | Func.TypeGuard _ ->
            let t = VoidT.at loc in
            let use_op =
              Op
                (FunImplicitReturn
                   { fn = reason_fn; upper = reason_of_t return_t; type_guard = true }
                )
            in
            (use_op, t, None)
        in

        let exhaust =
          match body with
          | None ->
            Flow.flow cx (void_t, UseT (use_op, return_t));
            None
          | Some _ ->
            let (exhaustive, undeclared) =
              try Context.exhaustive_check cx body_loc with
              | Not_found ->
                Flow_js_utils.add_output
                  cx
                  Error_message.(EInternal (loc, MissingSwitchExhaustiveCheck));
                ([], false)
            in
            Some
              ( ref
                  ( if undeclared then
                    [VoidT.at body_loc]
                  else
                    []
                  ),
                exhaustive,
                (void_t, UseT (use_op, return_t))
              )
        in
        (init_ast, exhaust)
      else
        (None, None)
    in

    let () =
      (new func_scope_visitor
         cx
         ~has_return_annot
         ~ret_annot_loc
         ~return_t
         ~yield_t
         ~next_t
         ~body_loc
         kind
         exhaust
      )
        #visit
        statements_ast
    in

    Base.Option.iter exhaust ~f:(fun (maybe_exhaustively_checked_ts, _, flow_on_non_exhaustive) ->
        if Type_operation_utils.TypeAssertions.non_exhaustive cx !maybe_exhaustively_checked_ts then
          Flow.flow cx flow_on_non_exhaustive
    );

    ignore @@ Type_env.set_scope_kind cx prev_scope_kind;

    (* return a tuple of (function body AST option, field initializer AST option).
       - the function body option is Some _ if the Param sig's body was Some, and
         None if the Param sig's body was None.
       - the field initializer is Some expr' if the Param sig's kind was FieldInit expr,
         where expr' is the typed AST translation of expr.
    *)
    (params_ast, body_ast, init_ast)

  let to_ctor_sig f = { f with T.kind = Func.Ctor }
end

let return_loc = function
  | { Ast.Function.return = Ast.Function.ReturnAnnot.Available (_, (loc, _)); _ }
  | { Ast.Function.return = Ast.Function.ReturnAnnot.TypeGuard (loc, _); _ }
  | { Ast.Function.body = Ast.Function.BodyExpression (loc, _); _ } ->
    loc
  | { Ast.Function.body = Ast.Function.BodyBlock (loc, _); _ } ->
    loc |> ALoc.to_loc_exn |> Loc.char_before |> ALoc.of_loc
