Compiler: compact program Addr.t (#1968)

* Compiler: compact program Addr.t

* Compatibility with OCaml 4.13

Author: hhugo

compiler/lib/code.ml

@@ -19,6 +19,10 @@
  *)
 open! Stdlib
 
+let stats = Debug.find "stats"
+
+let times = Debug.find "times"
+
 module Addr = struct
   type t = int
 
@@ -827,6 +831,65 @@ let with_invariant = Debug.find "invariant"
 
 let check_defs = false
 
+let do_compact { blocks; start; free_pc = _ } =
+  let remap =
+    let max = fst (Addr.Map.max_binding blocks) in
+    let a = Array.make (max + 1) 0 in
+    let i = ref 0 in
+    Addr.Map.iter
+      (fun pc _ ->
+        a.(pc) <- !i;
+        incr i)
+      blocks;
+    a
+  in
+  let rewrite_cont remap (pc, args) = remap.(pc), args in
+  let rewrite remap block =
+    let body =
+      List.map block.body ~f:(function
+        | Let (x, Closure (params, cont, loc)) ->
+            Let (x, Closure (params, rewrite_cont remap cont, loc))
+        | i -> i)
+    in
+    let branch =
+      match block.branch with
+      | (Return _ | Raise _ | Stop) as b -> b
+      | Branch c -> Branch (rewrite_cont remap c)
+      | Poptrap c -> Poptrap (rewrite_cont remap c)
+      | Cond (x, c1, c2) -> Cond (x, rewrite_cont remap c1, rewrite_cont remap c2)
+      | Switch (x, a) -> Switch (x, Array.map a ~f:(rewrite_cont remap))
+      | Pushtrap (c1, x, c2) -> Pushtrap (rewrite_cont remap c1, x, rewrite_cont remap c2)
+    in
+    { block with body; branch }
+  in
+  let blocks =
+    Addr.Map.fold
+      (fun pc b blocks -> Addr.Map.add remap.(pc) (rewrite remap b) blocks)
+      blocks
+      Addr.Map.empty
+  in
+  let free_pc = (Addr.Map.max_binding blocks |> fst) + 1 in
+  let start = remap.(start) in
+  { blocks; start; free_pc }
+
+let compact p =
+  let t = Timer.make () in
+  let card = Addr.Map.cardinal p.blocks in
+  let max = Addr.Map.max_binding p.blocks |> fst in
+  let ratio = float card /. float max *. 100. in
+  let do_it = Float.(ratio < 70.) in
+  let p = if do_it then do_compact p else p in
+  if times () then Format.eprintf "  compact: %a@." Timer.print t;
+  if stats ()
+  then
+    Format.eprintf
+      "Stats - compact: %d/%d = %.2f%%%s@."
+      card
+      max
+      ratio
+      (if not do_it then " - ignored" else "");
+  p
+
 let used_blocks p =
   let visited = BitSet.create' p.free_pc in
   let rec mark_used pc =

compiler/lib/code.mli

@@ -307,6 +307,8 @@ val prepend : program -> instr list -> program
 
 val empty : program
 
+val compact : program -> program
+
 val is_empty : program -> bool
 
 val equal : program -> program -> bool

compiler/lib/driver.ml

@@ -53,6 +53,7 @@ let deadcode' p =
 let deadcode p =
   let p, _ = deadcode' p in
   let p = Deadcode.merge_blocks p in
+  let p = Code.compact p in
   p
 
 let inline p =

compiler/lib/parse_bytecode.ml

@@ -2525,13 +2525,13 @@ let parse_bytecode code globals debug_data =
             })
           !compiled_blocks
       in
-      let free_pc = String.length code / 4 in
+      let free_pc = (Addr.Map.max_binding blocks |> fst) + 1 in
       { start; blocks; free_pc })
     else Code.empty
   in
   compiled_blocks := Addr.Map.empty;
   tagged_blocks := Addr.Map.empty;
-  p
+  Code.compact p
 
 module Toc : sig
   type t

compiler/lib/stdlib.ml

@@ -1100,6 +1100,18 @@ module In_channel = struct
 end
 [@@if ocaml_version >= (4, 14, 0)]
 
+module Seq = struct
+  include Seq
+
+  let rec mapi_aux f i xs () =
+    match xs () with
+    | Nil -> Nil
+    | Cons (x, xs) -> Cons (f i x, mapi_aux f (i + 1) xs)
+
+  (* Available since OCaml 4.14 *)
+  let[@inline] mapi f xs = mapi_aux f 0 xs
+end
+
 let split_lines s =
   if String.equal s ""
   then []