Source file globalize.ml

(* Wasm_of_ocaml compiler
 * http://www.ocsigen.org/js_of_ocaml/
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, with linking exception;
 * either version 2.1 of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *)

(*

Store some toplevel values into globals. Any variable which is used a
number of instructions after being defined is stored into a global
instead of a local. The goals are the following:
- Turn a large number of closures into constant closures, which has a
  significant impact on performance
- Reduce the compilation time of the toplevel function in case the
  Wasm engine decide to optimize it: reduce the register pressure by
  avoiding long-lived registers in the toplevel function, and make
  load elimination less expensive by reducing the number of constant
  structures defined in this function.
*)

open Stdlib

type st =
  { pos : int
  ; visited_variables : int Code.Var.Map.t
  ; globals : Code.Var.Set.t
  ; closures : Closure_conversion.closure Code.Var.Map.t
  ; constants : Code.Var.Set.t
  }

let threshold = 1000

let rec globalize st x =
  if Code.Var.Set.mem x st.globals
  then st
  else
    let st = { st with globals = Code.Var.Set.add x st.globals } in
    globalize_closure st x

and globalize_closure st x =
  (* If a function is stored in a global variable, its free variables
     are also stored in a global variable, since they are retained
     anyway. *)
  match Code.Var.Map.find x st.closures with
  | { free_variables; _ } ->
      List.fold_left
        ~f:(fun st x ->
          if Code.Var.Map.mem x st.visited_variables then globalize st x else st)
        ~init:st
        free_variables
  | exception Not_found -> st

let use x st =
  match Code.Var.Map.find x st.visited_variables with
  | pos -> if st.pos > pos + threshold then globalize st x else st
  | exception Not_found -> st

let declare x st =
  { st with visited_variables = Code.Var.Map.add x st.pos st.visited_variables }

let traverse_expression x e st =
  match e with
  | Code.Apply { f; args; _ } ->
      st |> use f |> fun st -> List.fold_left ~f:(fun st x -> use x st) ~init:st args
  | Block (_, a, _, _) -> Array.fold_right ~f:use a ~init:st
  | Field (x, _, _) -> st |> use x
  | Closure _ ->
      List.fold_left
        ~f:(fun st x -> use x st)
        ~init:st
        (Code.Var.Map.find x st.closures).Closure_conversion.free_variables
  | Constant _ -> { st with constants = Code.Var.Set.add x st.constants }
  | Special _ -> st
  | Prim (_, args) ->
      List.fold_left
        ~f:(fun st a ->
          match a with
          | Code.Pv x -> st |> use x
          | Pc _ -> st)
        ~init:st
        args

let traverse_instruction st i =
  let st = { st with pos = st.pos + 1 } in
  match i with
  | Code.Let (x, e) -> st |> declare x |> traverse_expression x e
  | Assign (_, x) | Offset_ref (x, _) -> st |> use x
  | Set_field (x, _, _, y) -> st |> use x |> use y
  | Array_set (x, y, z) -> st |> use x |> use y |> use z
  | Event _ -> st

let traverse_block p st pc =
  let b = Code.Addr.Map.find pc p.Code.blocks in
  let st = List.fold_left ~f:(fun st x -> declare x st) ~init:st b.Code.params in
  List.fold_left ~f:(fun st i -> traverse_instruction st i) ~init:st b.Code.body

let available x st = Code.Var.Set.mem x st.globals || Code.Var.Set.mem x st.constants

let propagate_instruction st i =
  match i with
  | Code.Let (x, Block (_, a, _, _)) when not (Code.Var.Set.mem x st.globals) ->
      (* Globalize a block when most of its fields are available
         (global or constant). Available fields go into the global's
         initializer; the rest are patched via [array.set] in the
         function body. The [+1] keeps 2-field cons cells eligible
         (one non-available field is allowed), which matters for
         cascading: globalizing an inner block makes its variable
         available for outer blocks. *)
      let non_available =
        Array.fold_right ~f:(fun v n -> if available v st then n else n + 1) a ~init:0
      in
      if 3 * non_available <= Array.length a + 1 then globalize st x else st
  | Code.Let (x, Closure _) when not (Code.Var.Set.mem x st.globals) -> (
      match Code.Var.Map.find x st.closures with
      | { free_variables; _ } ->
          if List.for_all ~f:(fun v -> available v st) free_variables
          then globalize st x
          else st
      | exception Not_found -> st)
  | _ -> st

let propagate_block p st pc =
  let b = Code.Addr.Map.find pc p.Code.blocks in
  List.fold_left ~f:(fun st i -> propagate_instruction st i) ~init:st b.Code.body

let f p g closures =
  let l = Structure.blocks_in_reverse_post_order g in
  let in_loop = Freevars.find_loops_in_closure p p.Code.start in
  let st =
    List.fold_left
      ~f:(fun st pc ->
        if Code.Addr.Map.mem pc in_loop then st else traverse_block p st pc)
      ~init:
        { pos = 0
        ; visited_variables = Code.Var.Map.empty
        ; globals = Code.Var.Set.empty
        ; closures
        ; constants = Code.Var.Set.empty
        }
      l
  in
  let st =
    List.fold_left
      ~f:(fun st pc ->
        if Code.Addr.Map.mem pc in_loop then st else propagate_block p st pc)
      ~init:st
      l
  in
  st.globals