Improved compilation of switches

Author: vouillon

compiler/lib/driver.ml

@@ -693,7 +693,8 @@ let optimize ~profile p =
     Code.Var.fresh_n "dummy"
   in
   let opt =
-    specialize_js_once
+    Specialize.switches
+    +> specialize_js_once
     +> (match profile with
        | O1 -> o1
        | O2 -> o2

compiler/lib/generate.ml

@@ -1120,7 +1120,7 @@ let _ =
   register_un_prim_ctx "%caml_format_int_special" `Pure (fun ctx cx loc ->
       let s = J.EBin (J.Plus, str_js_utf8 "", cx) in
       ocaml_string ~ctx ~loc s);
-  register_un_prim "%direct_obj_tag" `Mutator (fun cx _loc -> Mlvalue.Block.tag cx);
+  register_un_prim "%direct_obj_tag" `Pure (fun cx _loc -> Mlvalue.Block.tag cx);
   register_bin_prim "caml_array_unsafe_get" `Mutable (fun cx cy _ ->
       Mlvalue.Array.field cx cy);
   register_bin_prim "%int_add" `Pure (fun cx cy _ ->

compiler/lib/specialize.ml

@@ -117,3 +117,60 @@ let specialize_instrs ~function_arity p =
   { p with blocks; free_pc }
 
 let f = specialize_instrs
+
+(***)
+
+(* For switches, at this point, we know that this it is sufficient to
+   check the [pc]. *)
+let equal (pc, _) (pc', _) = pc = pc'
+
+let find_outlier_index arr =
+  let len = Array.length arr in
+  if len < 3
+  then if len = 1 || equal arr.(0) arr.(1) then `All_equals else `Distinguished 0
+  else
+    let majority =
+      if equal arr.(0) arr.(1) || equal arr.(0) arr.(2) then arr.(0) else arr.(1)
+    in
+    let rec find i =
+      if i >= len
+      then `All_equals
+      else if equal arr.(i) majority
+      then find (i + 1)
+      else `Distinguished i
+    in
+    match find 0 with
+    | `All_equals as res -> res
+    | `Distinguished i as res -> (
+        match find (i + 1) with
+        | `All_equals -> res
+        | `Distinguished _ -> `Many_cases)
+
+let switches p =
+  { p with
+    blocks =
+      Addr.Map.fold
+        (fun pc block blocks ->
+          match block.branch with
+          | Switch (x, l) -> (
+              match find_outlier_index l with
+              | `All_equals -> Addr.Map.add pc { block with branch = Branch l.(0) } blocks
+              | `Distinguished i ->
+                  let block =
+                    let v = Var.fresh () in
+                    let c = Var.fresh () in
+                    { block with
+                      body =
+                        block.body
+                        @ [ Let (v, Constant (Int (Targetint.of_int_exn i)))
+                          ; Let (c, Prim (Eq, [ Pv x; Pv v ]))
+                          ]
+                    ; branch = Cond (c, l.(i), l.((i + 1) mod Array.length l))
+                    }
+                  in
+                  Addr.Map.add pc block blocks
+              | `Many_cases -> blocks)
+          | _ -> blocks)
+        p.blocks
+        p.blocks
+  }

compiler/lib/specialize.mli

@@ -21,3 +21,5 @@
 val function_arity : Flow.Info.t -> Code.Var.t -> int option
 
 val f : function_arity:(Code.Var.t -> int option) -> Code.program -> Code.program
+
+val switches : Code.program -> Code.program