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

open Flow_ast_visitor
open Hoister
open Scope_api

module LocMap = Utils_js.LocMap

class with_or_eval_visitor = object(this)
  inherit [bool] visitor ~init:false as super

  method! expression (expr: Loc.t Ast.Expression.t) =
    let open Ast.Expression in
    if this#acc = true then expr else match expr with
    | (_, Call { Call.callee = (_, Identifier (_, "eval")); _}) ->
      this#set_acc true;
      expr
    | _ -> super#expression expr

  method! statement (stmt: Loc.t Ast.Statement.t) =
    if this#acc = true then stmt else super#statement stmt

  method! with_ (stuff: Loc.t Ast.Statement.With.t) =
    this#set_acc true;
    stuff
end

(* Visitor class that prepares use-def info, hoisting bindings one scope at a
   time. This info can be used for various purposes, e.g. variable renaming.

   We do not generate the scope tree for the entire program, because it is not
   clear where to hang scopes for function expressions, catch clauses,
   etc. One possibility is to augment the AST with scope identifiers.

   As we move into a nested scope, we generate bindings for the new scope, map
   the bindings to names generated by a factory, and augment the existing
   environment with this map before visiting the nested scope.

   Because globals can appear deep in the program, we cannot actually perform
   any renaming until we have walked the entire program. Instead, we compute (1)
   a map from identifier locations to name ids corresponding to binding sites
   and (2) a map from name ids to globals they conflict with. Later, we generate
   names (avoiding conflicts with globals) and rename those locations.

   We try to create the smallest number of name ids possible.
*)
module Acc = struct
  type t = info
  let init = {
    max_distinct = 0;
    globals = IMap.empty;
    locals = LocMap.empty;
    scopes = IMap.empty;
  }
end

class scope_builder = object(this)
  inherit [Acc.t] visitor ~init:Acc.init as super

  val mutable env = SMap.empty
  val mutable current_scope_opt = None
  val mutable scope_counter = 0
  method private new_scope scope =
    let new_scope = scope_counter in
    scope_counter <- scope_counter + 1;
    this#update_acc (fun acc -> { acc with
      scopes = IMap.add new_scope scope acc.scopes
    });
    new_scope
  method current_scope =
    match current_scope_opt with
    | None -> failwith "global scope"
    | Some i -> i

  val mutable counter = 0
  method private next =
    let result = counter in
    counter <- counter + 1;
    this#update_acc (fun acc -> { acc with
      max_distinct = max counter acc.max_distinct
    });
    result

  method private mk_env parent_env scope =
    List.fold_left (fun map (x, locs) ->
      let def = match SMap.get x parent_env with
        | Some def -> Def.{ def with locs; scope; }
        | None -> Def.{ locs; scope; name = this#next; } in
      SMap.add x def map
    ) SMap.empty

  method private push ?(lexical=false) bindings =
    let save_counter = counter in
    let old_env = env in
    let parent = current_scope_opt in
    let child = this#new_scope Scope.{ lexical; parent; } in
    current_scope_opt <- Some child;
    env <- SMap.fold SMap.add (this#mk_env old_env child (Bindings.to_assoc bindings)) old_env;
    parent, old_env, save_counter

  method private pop (parent, old_env, save_counter) =
    current_scope_opt <- parent;
    env <- old_env;
    counter <- save_counter

  method with_bindings: 'a. ?lexical:bool -> Bindings.t -> ('a -> 'a) -> 'a -> 'a =
    fun ?lexical bindings visit node ->
      let saved_state = this#push ?lexical bindings in
      let node' = visit node in
      this#pop saved_state;
      node'

  method private add_global x =
    SMap.iter (fun _ { Def.name; _ } ->
      this#update_acc (fun acc -> { acc with
        globals =
          let iglobals = try IMap.find_unsafe name acc.globals with _ -> SSet.empty in
          IMap.add name (SSet.add x iglobals) acc.globals
      })
    ) env

  method private add_local loc def =
    this#update_acc (fun acc -> { acc with
      locals = LocMap.add loc (def, this#current_scope) acc.locals
    })

  method! identifier (expr: Loc.t Ast.Identifier.t) =
    let loc, x = expr in
    begin match SMap.get x env with
      | Some def -> this#add_local loc def
      | None -> this#add_global x
    end;
    expr

  (* don't rename the `foo` in `x.foo` *)
  method! member_property_identifier (id: Loc.t Ast.Identifier.t) = id

  (* don't rename the `foo` in `{ foo: ... }` *)
  method! object_key_identifier (id: Loc.t Ast.Identifier.t) = id

  method! block (stmt: Loc.t Ast.Statement.Block.t) =
    let lexical_hoist = new lexical_hoister in
    let lexical_bindings = lexical_hoist#eval lexical_hoist#block stmt in
    this#with_bindings ~lexical:true lexical_bindings super#block stmt

  (* like block *)
  method! program (program: Loc.t Ast.program) =
    let lexical_hoist = new lexical_hoister in
    let lexical_bindings = lexical_hoist#eval lexical_hoist#program program in
    this#with_bindings ~lexical:true lexical_bindings super#program program

  method private scoped_for_in_statement (stmt: Loc.t Ast.Statement.ForIn.t) =
    super#for_in_statement stmt

  method! for_in_statement (stmt: Loc.t Ast.Statement.ForIn.t) =
    let open Ast.Statement.ForIn in
    let { left; right = _; body = _; each = _ } = stmt in

    let lexical_hoist = new lexical_hoister in
    let lexical_bindings = match left with
    | LeftDeclaration (_, decl) ->
      lexical_hoist#eval lexical_hoist#variable_declaration decl
    | LeftPattern _ -> Bindings.empty
    in
    this#with_bindings ~lexical:true lexical_bindings this#scoped_for_in_statement stmt

  method private scoped_for_of_statement (stmt: Loc.t Ast.Statement.ForOf.t) =
    super#for_of_statement stmt

  method! for_of_statement (stmt: Loc.t Ast.Statement.ForOf.t) =
    let open Ast.Statement.ForOf in
    let { left; right = _; body = _; async = _ } = stmt in

    let lexical_hoist = new lexical_hoister in
    let lexical_bindings = match left with
    | LeftDeclaration (_, decl) ->
      lexical_hoist#eval lexical_hoist#variable_declaration decl
    | LeftPattern _ -> Bindings.empty
    in
    this#with_bindings ~lexical:true lexical_bindings this#scoped_for_of_statement stmt

  method private scoped_for_statement (stmt: Loc.t Ast.Statement.For.t) =
    super#for_statement stmt

  method! for_statement (stmt: Loc.t Ast.Statement.For.t) =
    let open Ast.Statement.For in
    let { init; test = _; update = _; body = _ } = stmt in

    let lexical_hoist = new lexical_hoister in
    let lexical_bindings = match init with
    | Some (InitDeclaration (_, decl)) ->
      lexical_hoist#eval lexical_hoist#variable_declaration decl
    | _ -> Bindings.empty
    in
    this#with_bindings ~lexical:true lexical_bindings this#scoped_for_statement stmt

  method! catch_clause (clause: Loc.t Ast.Statement.Try.CatchClause.t') =
    let open Ast.Statement.Try.CatchClause in
    let { param; body = _ } = clause in

    this#with_bindings (
      let open Ast.Pattern in
      let _, patt = param in
      match patt with
      | Identifier { Identifier.name; _ } -> Bindings.singleton name
      | _ -> (* TODO *)
        Bindings.empty
    ) super#catch_clause clause

  (* helper for function params and body *)
  method private lambda params body =
    let open Ast.Function in

    (* hoisting *)
    let hoist = new hoister in
    begin
      let (_loc, { Params.params = param_list; rest = _rest }) = params in
      run_list hoist#function_param_pattern param_list;
      match body with
        | BodyBlock (_loc, block) ->
          run hoist#block block
        | _ ->
          ()
    end;

    this#with_bindings hoist#acc (fun () ->
      let (_loc, { Params.params = param_list; rest }) = params in
      run_list this#function_param_pattern param_list;
      run_opt this#function_rest_element rest;
      begin match body with
      | BodyBlock (_, block) ->
        run this#block block
      | BodyExpression expr ->
        run this#expression expr
      end;
    ) ()

  method! function_declaration (expr: Loc.t Ast.Function.t) =
    let contains_with_or_eval =
      let visit = new with_or_eval_visitor in
      visit#eval visit#function_declaration expr
    in

    if not contains_with_or_eval then begin
      let open Ast.Function in
      let {
        id; params; body; async = _; generator = _; expression = _;
        predicate = _; returnType = _; typeParameters = _;
      } = expr in

      run_opt this#function_identifier id;

      this#lambda params body;
    end;

    expr

  (* Almost the same as function_declaration, except that the name of the
     function expression is locally in scope. *)
  method! function_ (expr: Loc.t Ast.Function.t) =
    let contains_with_or_eval =
      let visit = new with_or_eval_visitor in
      visit#eval visit#function_ expr
    in

    if not contains_with_or_eval then begin
      let open Ast.Function in
      let {
        id; params; body; async = _; generator = _; expression = _;
        predicate = _; returnType = _; typeParameters = _;
      } = expr in

      let bindings = match id with
        | Some name -> Bindings.singleton name
        | None -> Bindings.empty in
      this#with_bindings bindings (fun () ->
        run_opt this#function_identifier id;
        this#lambda params body;
      ) ();
    end;

    expr
end

let program ?(ignore_toplevel=false) program =
  let walk = new scope_builder in
  if ignore_toplevel then walk#eval walk#program program
  else
    let hoist = new hoister in
    let bindings = hoist#eval hoist#program program in
    walk#eval (walk#with_bindings bindings walk#program) program
