(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*                       Pierre Chambart, OCamlPro                        *)
(*           Mark Shinwell and Leo White, Jane Street Europe              *)
(*                                                                        *)
(*   Copyright 2013--2016 OCamlPro SAS                                    *)
(*   Copyright 2014--2016 Jane Street Group LLC                           *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

[@@@ocaml.warning "+a-4-9-30-40-41-42-66"]
open! Int_replace_polymorphic_compare

module B = Inlining_cost.Benefit

let pass_name = "unbox-free-vars-of-closures"
let () = Pass_wrapper.register ~pass_name

(* CR-someday mshinwell: Nearly but not quite the same as something that
   Augment_specialised_args uses. *)
let add_lifted_projections_around_set_of_closures
      ~set_of_closures ~existing_inner_to_outer_vars ~benefit
      ~definitions_indexed_by_new_inner_vars =
  let body =
    Flambda_utils.name_expr (Set_of_closures set_of_closures)
      ~name:Internal_variable_names.unbox_free_vars_of_closures
  in
  Variable.Map.fold (fun new_inner_var (projection : Projection.t)
            (expr, benefit) ->
      let find_outer_var inner_var =
        match
          Variable.Map.find inner_var existing_inner_to_outer_vars
        with
        | (outer_var : Flambda.specialised_to) -> outer_var.var
        | exception Not_found ->
          Misc.fatal_errorf "(UFV) find_outer_var: expected %a \
              to be in [existing_inner_to_outer_vars], but it is \
              not.  (The projection was: %a)"
            Variable.print inner_var
            Projection.print projection
      in
      let benefit = B.add_projection projection benefit in
      let named : Flambda.named =
        (* The lifted projection must be in terms of outer variables,
           not inner variables. *)
        let projection =
          Projection.map_projecting_from projection ~f:find_outer_var
        in
        Flambda_utils.projection_to_named projection
      in
      let expr =
        Flambda.create_let (find_outer_var new_inner_var) named expr
      in
      (expr, benefit))
    definitions_indexed_by_new_inner_vars
    (body, benefit)

let run ~env ~(set_of_closures : Flambda.set_of_closures) =
  if not !Clflags.unbox_free_vars_of_closures then
    None
  else
    let definitions_indexed_by_new_inner_vars, _, free_vars, done_something =
      let all_existing_definitions =
        Variable.Map.fold (fun _inner_var (outer_var : Flambda.specialised_to)
              all_existing_definitions ->
            match outer_var.projection with
            | None -> all_existing_definitions
            | Some projection ->
              Projection.Set.add projection all_existing_definitions)
          set_of_closures.free_vars
          Projection.Set.empty
      in
      Flambda_iterators.fold_function_decls_ignoring_stubs set_of_closures
        ~init:(Variable.Map.empty, all_existing_definitions,
          set_of_closures.free_vars, false)
        ~f:(fun ~fun_var:_ ~function_decl result ->
          let extracted =
            Extract_projections.from_function_decl ~env ~function_decl
              ~which_variables:set_of_closures.free_vars
          in
          Projection.Set.fold (fun projection
                ((definitions_indexed_by_new_inner_vars,
                  all_existing_definitions_including_added_ones,
                  additional_free_vars, _done_something) as result) ->
              (* Don't add a new free variable if there already exists a
                 free variable with the desired projection.  We need to
                 dedup not only across the existing free variables but
                 also across newly-added ones (unlike in
                 [Augment_specialised_args]), since free variables are
                 not local to a function declaration but rather to a
                 set of closures. *)
              if Projection.Set.mem projection
                all_existing_definitions_including_added_ones
              then begin
                result
              end else begin
                (* Add a new free variable.  This needs both a fresh
                   "new inner" and a fresh "new outer" var, since we know
                   the definition is not a duplicate. *)
                let projecting_from = Projection.projecting_from projection in
                let new_inner_var = Variable.rename projecting_from in
                let new_outer_var = Variable.rename projecting_from in
                let definitions_indexed_by_new_inner_vars =
                  Variable.Map.add new_inner_var projection
                    definitions_indexed_by_new_inner_vars
                in
                let all_existing_definitions_including_added_ones =
                  Projection.Set.add projection
                    all_existing_definitions_including_added_ones
                in
                let new_outer_var : Flambda.specialised_to =
                  { var = new_outer_var;
                    projection = Some projection;
                  }
                in
                let additional_free_vars =
                  Variable.Map.add new_inner_var new_outer_var
                    additional_free_vars
                in
                definitions_indexed_by_new_inner_vars,
                  all_existing_definitions_including_added_ones,
                  additional_free_vars,
                  true
              end)
            extracted
            result)
    in
    if not done_something then
      None
    else
      (* CR-someday mshinwell: could consider doing the grouping thing
         similar to Augment_specialised_args *)
      let num_free_vars_before =
        Variable.Map.cardinal set_of_closures.free_vars
      in
      let num_free_vars_after =
        Variable.Map.cardinal free_vars
      in
      assert (num_free_vars_after > num_free_vars_before);
      (* Don't let the closure grow too large. *)
      if num_free_vars_after > 2 * num_free_vars_before then
        None
      else
        let set_of_closures =
          Flambda.create_set_of_closures
            ~function_decls:set_of_closures.function_decls
            ~free_vars
            ~specialised_args:set_of_closures.specialised_args
            ~direct_call_surrogates:set_of_closures.direct_call_surrogates
        in
        let expr, benefit =
          add_lifted_projections_around_set_of_closures ~set_of_closures
            ~benefit:B.zero
            ~existing_inner_to_outer_vars:set_of_closures.free_vars
            ~definitions_indexed_by_new_inner_vars
        in
        Some (expr, benefit)

let run ~env ~set_of_closures =
  Pass_wrapper.with_dump ~ppf_dump:(Inline_and_simplify_aux.Env.ppf_dump env)
    ~pass_name ~input:set_of_closures
    ~print_input:Flambda.print_set_of_closures
    ~print_output:(fun ppf (expr, _) -> Flambda.print ppf expr)
    ~f:(fun () -> run ~env ~set_of_closures)
