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

module Make (Env : Env_sig.S) : Abnormal_sig.S with module Env := Env = struct
  (* we model abnormal control flows using exceptions during traversal *)

  (* control directives encountered during traversal *)
  type t =
    | Return
    | Throw
    | Break of string option
    | Continue of string option

  type payload =
    | Expr of ALoc.t * (ALoc.t, ALoc.t * Type.t) Flow_ast.Expression.t
    | Stmt of (ALoc.t, ALoc.t * Type.t) Flow_ast.Statement.t
    | Stmts of (ALoc.t, ALoc.t * Type.t) Flow_ast.Statement.t list

  exception Exn of payload * t

  open Utils_js

  (* called from traversal. abnormal indicates control flow directive encountered *)
  let throw_stmt_control_flow_exception stmt abnormal = raise (Exn (Stmt stmt, abnormal))

  let throw_stmts_control_flow_exception stmts abnormal = raise (Exn (Stmts stmts, abnormal))

  let throw_expr_control_flow_exception loc expr abnormal = raise (Exn (Expr (loc, expr), abnormal))

  (* if argument is Some abnormal, throw it *)
  let check_stmt_control_flow_exception = function
    | (stmt, None) -> stmt
    | (stmt, Some abnormal) -> throw_stmt_control_flow_exception stmt abnormal

  (* helper *)
  let check_env_depth depth =
    let new_depth = Env.env_depth () in
    if new_depth = depth then
      ()
    else
      assert_false (spf "env depth %d != %d after no control flow catch" new_depth depth)

  (* catch_stmt_control_flow_exception runs a function which is expected to either
       return a statement or raise Exn (Stmt _, _) or Exn (Expr _, _). The function
       should never raise Exn (Stmts _, _). If the function raises an
       Exn (Expr _, _), an Expression statement will be constructed around the
       expression.
     Similarly, the function passed into catch_stmt_control_flow_exception should
       return a statement list or raise Exn (Stmts _, _), and never raise
       Exn (Stmt _, _) or Exn (Expr _, _); and likewise
       catch_expr_control_flow_exception should raise Expr and not Stmts or Stmt
     For all:
       If the passed-in function returns an AST,
         then we return that AST and None.
       Otherwise, if it raises with some AST payload and an abnormal flow,
         then we return the payload AST and Some <abnormal flow>.
  *)
  let ( catch_stmt_control_flow_exception,
        catch_stmts_control_flow_exception,
        catch_expr_control_flow_exception
      ) =
    let catch_control_flow_exception p f =
      let depth = Env.env_depth () in
      try
        let res = f () in
        check_env_depth depth;
        (res, None)
      with
      | Exn (payload, abnormal) ->
        Env.trunc_env depth;
        (p payload, Some abnormal)
      | exn -> raise exn
    in
    ( catch_control_flow_exception (function
          | Stmt stmt -> stmt
          (* If we catch an Expr payload, then it was generated by seeing an invariant() call
             and we can reconstruct a statement around it (rather than constructing a statement
             around it when we create the payload) *)
          | Expr (loc, exp) ->
            ( loc,
              Flow_ast.Statement.Expression
                {
                  Flow_ast.Statement.Expression.expression = exp;
                  directive = None;
                  comments = None;
                }
            )
          | Stmts _ -> assert_false "Statement expected"
          ),
      catch_control_flow_exception (function
          | Stmts stmts -> stmts
          | Stmt _
          | Expr _ ->
            assert_false "Statement list expected"
          ),
      catch_control_flow_exception (function
          | Expr (_, exp) -> exp
          | Stmt _
          | Stmts _ ->
            assert_false "Expression expected"
          )
    )

  (* like check_control_flow_exception, except break statements
     specifying the given label (or None) are ignored *)
  let ignore_break_to_label label = function
    | (ast, Some (Break break_label)) when break_label = label -> (ast, None)
    | result -> result

  (* like ignore_break_to_label, except continue statements
     on the same label (or None) are also ignored *)
  let ignore_break_or_continue_to_label label res =
    match ignore_break_to_label label res with
    | (ast, Some (Continue cont_label)) when cont_label = label -> (ast, None)
    | result -> result

  (********************************************************************)

  (** at some points we need to record control flow directives in addition
    to responding to them. *)

  module AbnormalMap : WrappedMap.S with type key = t = WrappedMap.Make (struct
    type abnormal = t

    type t = abnormal

    let compare = Stdlib.compare
  end)

  let abnormals : Env.t AbnormalMap.t ref = ref AbnormalMap.empty

  (** record the appearance of a control flow directive.
    associate the given env if passed *)
  let save ?(env = []) abnormal = abnormals := AbnormalMap.add abnormal env !abnormals

  (** set or remove a given control flow directive's value,
    and return the current one *)
  let swap_saved abnormal value =
    let old = AbnormalMap.find_opt abnormal !abnormals in
    ( if old <> value then
      abnormals :=
        match value with
        | None -> AbnormalMap.remove abnormal !abnormals
        | Some env -> AbnormalMap.add abnormal env !abnormals
    );
    old

  (** remove a given control flow directive's value,
    and return the current one *)
  let clear_saved abnormal = swap_saved abnormal None

  let try_with_abnormal_exn ~f ~on_abnormal_exn () =
    try f () with
    | Exn (payload, t) -> on_abnormal_exn (payload, t)
    | exc -> raise exc
end
