Compiler: lower level switch

Author: hhugo

compiler/lib/code.ml

@@ -354,7 +354,7 @@ type last =
   | Stop
   | Branch of cont
   | Cond of Var.t * cont * cont
-  | Switch of Var.t * cont array * cont array
+  | Switch of Var.t * cont array
   | Pushtrap of cont * Var.t * cont * Addr.Set.t
   | Poptrap of cont
 
@@ -499,10 +499,9 @@ module Print = struct
     | Branch c -> Format.fprintf f "branch %a" cont c
     | Cond (x, cont1, cont2) ->
         Format.fprintf f "if %a then %a else %a" Var.print x cont cont1 cont cont2
-    | Switch (x, a1, a2) ->
+    | Switch (x, a1) ->
         Format.fprintf f "switch %a {" Var.print x;
         Array.iteri a1 ~f:(fun i c -> Format.fprintf f "int %d -> %a; " i cont c);
-        Array.iteri a2 ~f:(fun i c -> Format.fprintf f "tag %d -> %a; " i cont c);
         Format.fprintf f "}"
     | Pushtrap (cont1, x, cont2, pcs) ->
         Format.fprintf
@@ -598,9 +597,8 @@ let fold_children blocks pc f accu =
       let accu = f pc1 accu in
       let accu = f pc2 accu in
       accu
-  | Switch (_, a1, a2) ->
+  | Switch (_, a1) ->
       let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a1 in
-      let accu = Array.fold_right ~init:accu ~f:(fun (pc, _) accu -> f pc accu) a2 in
       accu
 
 type 'c fold_blocs = block Addr.Map.t -> Addr.t -> (Addr.t -> 'c -> 'c) -> 'c -> 'c
@@ -726,9 +724,7 @@ let invariant { blocks; start; _ } =
       | Cond (_x, cont1, cont2) ->
           check_cont cont1;
           check_cont cont2
-      | Switch (_x, a1, a2) ->
-          Array.iteri a1 ~f:(fun _ cont -> check_cont cont);
-          Array.iteri a2 ~f:(fun _ cont -> check_cont cont)
+      | Switch (_x, a1) -> Array.iteri a1 ~f:(fun _ cont -> check_cont cont)
       | Pushtrap (cont1, _x, cont2, _pcs) ->
           check_cont cont1;
           check_cont cont2

compiler/lib/code.mli

@@ -201,7 +201,7 @@ type last =
   | Stop
   | Branch of cont
   | Cond of Var.t * cont * cont
-  | Switch of Var.t * cont array * cont array
+  | Switch of Var.t * cont array
   | Pushtrap of cont * Var.t * cont * Addr.Set.t
   | Poptrap of cont
 

compiler/lib/deadcode.ml

@@ -97,10 +97,9 @@ and mark_reachable st pc =
         mark_var st x;
         mark_cont_reachable st cont1;
         mark_cont_reachable st cont2
-    | Switch (x, a1, a2) ->
+    | Switch (x, a1) ->
         mark_var st x;
-        Array.iter a1 ~f:(fun cont -> mark_cont_reachable st cont);
-        Array.iter a2 ~f:(fun cont -> mark_cont_reachable st cont)
+        Array.iter a1 ~f:(fun cont -> mark_cont_reachable st cont)
     | Pushtrap (cont1, _, cont2, _) ->
         mark_cont_reachable st cont1;
         mark_cont_reachable st cont2)
@@ -136,11 +135,8 @@ let filter_live_last blocks st (l, loc) =
     | Branch cont -> Branch (filter_cont blocks st cont)
     | Cond (x, cont1, cont2) ->
         Cond (x, filter_cont blocks st cont1, filter_cont blocks st cont2)
-    | Switch (x, a1, a2) ->
-        Switch
-          ( x
-          , Array.map a1 ~f:(fun cont -> filter_cont blocks st cont)
-          , Array.map a2 ~f:(fun cont -> filter_cont blocks st cont) )
+    | Switch (x, a1) ->
+        Switch (x, Array.map a1 ~f:(fun cont -> filter_cont blocks st cont))
     | Pushtrap (cont1, x, cont2, pcs) ->
         Pushtrap
           ( filter_cont blocks st cont1
@@ -204,9 +200,7 @@ let f ({ blocks; _ } as p : Code.program) =
       | Cond (_, cont1, cont2) ->
           add_cont_dep blocks defs cont1;
           add_cont_dep blocks defs cont2
-      | Switch (_, a1, a2) ->
-          Array.iter a1 ~f:(fun cont -> add_cont_dep blocks defs cont);
-          Array.iter a2 ~f:(fun cont -> add_cont_dep blocks defs cont)
+      | Switch (_, a1) -> Array.iter a1 ~f:(fun cont -> add_cont_dep blocks defs cont)
       | Pushtrap (cont, _, cont_h, _) ->
           add_cont_dep blocks defs cont_h;
           add_cont_dep blocks defs cont

compiler/lib/effects.ml

@@ -418,13 +418,11 @@ let cps_last ~st ~alloc_jump_closures pc ((last, last_loc) : last * loc) ~k :
             , cps_jump_cont ~st ~src:pc cont1 last_loc
             , cps_jump_cont ~st ~src:pc cont2 last_loc )
         , last_loc ) )
-  | Switch (x, c1, c2) ->
+  | Switch (x, c1) ->
       (* To avoid code duplication during JavaScript generation, we need
          to create a single block per continuation *)
       let cps_jump_cont = Fun.memoize (fun x -> cps_jump_cont ~st ~src:pc x last_loc) in
-      ( alloc_jump_closures
-      , ( Switch (x, Array.map c1 ~f:cps_jump_cont, Array.map c2 ~f:cps_jump_cont)
-        , last_loc ) )
+      alloc_jump_closures, (Switch (x, Array.map c1 ~f:cps_jump_cont), last_loc)
   | Pushtrap (body_cont, exn, ((handler_pc, _) as handler_cont), _) -> (
       assert (Hashtbl.mem st.is_continuation handler_pc);
       match Addr.Set.mem handler_pc st.blocks_to_transform with
@@ -911,8 +909,7 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
                match branch with
                | Branch cont -> Branch (resolve cont)
                | Cond (x, cont1, cont2) -> Cond (x, resolve cont1, resolve cont2)
-               | Switch (x, a1, a2) ->
-                   Switch (x, Array.map ~f:resolve a1, Array.map ~f:resolve a2)
+               | Switch (x, a1) -> Switch (x, Array.map ~f:resolve a1)
                | Pushtrap (cont1, x, cont2, s) ->
                    Pushtrap (resolve cont1, x, resolve cont2, s)
                | Poptrap cont -> Poptrap (resolve cont)

compiler/lib/eval.ml

@@ -341,14 +341,13 @@ let eval_branch info (l, loc) =
           | Zero -> Branch ffalse
           | Non_zero -> Branch ftrue
           | Unknown -> b)
-    | Switch (x, const, tags) as b -> (
+    | Switch (x, const) as b -> (
         (* [the_case_of info (Pv x)] might be meaningless when we're inside a dead code.
            The proper fix would be to remove the deadcode entirely.
            Meanwhile, add guards to prevent Invalid_argument("index out of bounds")
            see https://github.com/ocsigen/js_of_ocaml/issues/485 *)
         match the_case_of info (Pv x) with
         | CConst j when j >= 0 && j < Array.length const -> Branch const.(j)
-        | CTag j when j >= 0 && j < Array.length tags -> Branch tags.(j)
         | CConst _ | CTag _ | Unknown -> b)
     | _ as b -> b
   in
@@ -380,15 +379,11 @@ let rec do_not_raise pc visited blocks =
         let visited = do_not_raise pc1 visited blocks in
         let visited = do_not_raise pc2 visited blocks in
         visited
-    | Switch (_, a1, a2) ->
+    | Switch (_, a1) ->
         let visited =
           Array.fold_left a1 ~init:visited ~f:(fun visited (pc, _) ->
               do_not_raise pc visited blocks)
         in
-        let visited =
-          Array.fold_left a2 ~init:visited ~f:(fun visited (pc, _) ->
-              do_not_raise pc visited blocks)
-        in
         visited
     | Pushtrap _ -> raise May_raise
 

compiler/lib/flow.ml

@@ -119,9 +119,8 @@ let program_deps { blocks; _ } =
       | Cond (_, cont1, cont2) ->
           cont_deps blocks vars deps defs cont1;
           cont_deps blocks vars deps defs cont2
-      | Switch (_, a1, a2) ->
-          Array.iter a1 ~f:(fun cont -> cont_deps blocks vars deps defs cont);
-          Array.iter a2 ~f:(fun cont -> cont_deps blocks vars deps defs cont)
+      | Switch (_, a1) ->
+          Array.iter a1 ~f:(fun cont -> cont_deps blocks vars deps defs cont)
       | Pushtrap (cont, x, cont_h, _) ->
           add_param_def vars defs x;
           cont_deps blocks vars deps defs cont_h;

compiler/lib/freevars.ml

@@ -64,10 +64,9 @@ let iter_last_free_var f l =
       f x;
       iter_cont_free_vars f cont1;
       iter_cont_free_vars f cont2
-  | Switch (x, a1, a2) ->
+  | Switch (x, a1) ->
       f x;
-      Array.iter a1 ~f:(fun c -> iter_cont_free_vars f c);
-      Array.iter a2 ~f:(fun c -> iter_cont_free_vars f c)
+      Array.iter a1 ~f:(fun c -> iter_cont_free_vars f c)
   | Pushtrap (cont1, _, cont2, _) ->
       iter_cont_free_vars f cont1;
       iter_cont_free_vars f cont2

compiler/lib/generate.ml

@@ -760,10 +760,9 @@ let fold_children blocks pc f accu =
       let accu = f pc2 accu in
       accu
   | Cond (_, cont1, cont2) -> DTree.fold_cont f (DTree.build_if cont1 cont2) accu
-  | Switch (_, a1, a2) ->
-      let a1 = DTree.build_switch a1 and a2 = DTree.build_switch a2 in
+  | Switch (_, a1) ->
+      let a1 = DTree.build_switch a1 in
       let accu = DTree.fold_cont f a1 accu in
-      let accu = DTree.fold_cont f a2 accu in
       accu
 
 let build_graph ctx pc =
@@ -1099,6 +1098,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_bin_prim "caml_array_unsafe_get" `Mutable (fun cx cy _ ->
       Mlvalue.Array.field cx cy);
   register_bin_prim "%int_add" `Pure (fun cx cy _ -> to_int (plus_int cx cy));
@@ -1769,7 +1769,7 @@ and colapse_frontier name st (new_frontier' : Addr.Set.t) interm =
     let branch =
       let cases = Array.of_list (List.map a ~f:(fun pc -> pc, [])) in
       if Array.length cases > 2
-      then Code.Switch (x, cases, [||]), Code.noloc
+      then Code.Switch (x, cases), Code.noloc
       else Code.Cond (x, cases.(1), cases.(0)), Code.noloc
     in
     ( [ J.variable_declaration [ J.V x, (int default, J.N) ], J.N ]
@@ -1850,7 +1850,7 @@ and compile_conditional st queue last loop_stack backs frontier interm =
      | Raise _ -> Format.eprintf "raise;@;"
      | Stop -> Format.eprintf "stop;@;"
      | Cond (x, _, _) -> Format.eprintf "@[<hv 2>cond(%a){@;" Code.Var.print x
-     | Switch (x, _, _) -> Format.eprintf "@[<hv 2>switch(%a){@;" Code.Var.print x);
+     | Switch (x, _) -> Format.eprintf "@[<hv 2>switch(%a){@;" Code.Var.print x);
   let loc = source_location st.ctx pc in
   let res =
     match last with
@@ -1912,21 +1912,7 @@ and compile_conditional st queue last loop_stack backs frontier interm =
             (DTree.build_if c1 c2)
         in
         never, flush_all queue b
-    | Switch (x, [||], a2) ->
-        let (_px, cx), queue = access_queue queue x in
-        let never, code =
-          compile_decision_tree
-            st
-            loop_stack
-            backs
-            frontier
-            interm
-            loc
-            (Mlvalue.Block.tag cx)
-            (DTree.build_switch a2)
-        in
-        never, flush_all queue code
-    | Switch (x, a1, [||]) ->
+    | Switch (x, a1) ->
         let (_px, cx), queue = access_queue queue x in
         let never, code =
           compile_decision_tree
@@ -1940,41 +1926,6 @@ and compile_conditional st queue last loop_stack backs frontier interm =
             (DTree.build_switch a1)
         in
         never, flush_all queue code
-    | Switch (x, a1, a2) ->
-        (* The variable x is accessed several times, so we can directly
-           refer to it *)
-        let never1, b1 =
-          compile_decision_tree
-            st
-            loop_stack
-            backs
-            frontier
-            interm
-            loc
-            (var x)
-            (DTree.build_switch a1)
-        in
-        let never2, b2 =
-          compile_decision_tree
-            st
-            loop_stack
-            backs
-            frontier
-            interm
-            loc
-            (Mlvalue.Block.tag (var x))
-            (DTree.build_switch a2)
-        in
-        let code =
-          Js_simpl.if_statement
-            (Mlvalue.is_immediate (var x))
-            loc
-            (Js_simpl.block b1)
-            never1
-            (Js_simpl.block b2)
-            never2
-        in
-        never1 && never2, flush_all queue code
   in
   (if debug ()
    then

compiler/lib/global_flow.ml

@@ -243,14 +243,9 @@ let program_deps st { blocks; _ } =
       | Cond (x, cont1, cont2) ->
           cont_deps blocks st cont1;
           cont_deps blocks st ~ignore:x cont2
-      | Switch (x, a1, a2) ->
+      | Switch (x, a1) ->
           Array.iter a1 ~f:(fun cont -> cont_deps blocks st cont);
-          Array.iter a2 ~f:(fun cont -> cont_deps blocks st cont);
           let h = Hashtbl.create 16 in
-          Array.iteri
-            ~f:(fun i (pc, _) ->
-              Hashtbl.replace h pc (i :: (try Hashtbl.find h pc with Not_found -> [])))
-            a2;
           if not st.fast
           then
             Hashtbl.iter

compiler/lib/inline.ml

@@ -38,7 +38,7 @@ let optimizable blocks pc _ =
             +
             match fst branch with
             | Cond _ -> 2
-            | Switch (_, a1, a2) -> Array.length a1 + Array.length a2
+            | Switch (_, a1) -> Array.length a1
             | _ -> 0)
       in
       let optimizable =
@@ -112,9 +112,8 @@ let fold_children blocks pc f accu =
       let accu = f pc1 accu in
       let accu = f pc2 accu in
       accu
-  | Switch (_, a1, a2) ->
+  | Switch (_, a1) ->
       let accu = Array.fold_right a1 ~init:accu ~f:(fun (pc, _) accu -> f pc accu) in
-      let accu = Array.fold_right a2 ~init:accu ~f:(fun (pc, _) accu -> f pc accu) in
       accu
 
 let rewrite_closure blocks cont_pc clos_pc =

compiler/lib/parse_bytecode.ml

@@ -850,10 +850,9 @@ let rec compile_block blocks debug_data code pc state =
     | Cond (_, (pc1, _), (pc2, _)) ->
         compile_block blocks debug_data code pc1 state';
         compile_block blocks debug_data code pc2 state'
-    | Switch (_, l1, l2) ->
-        Array.iter l1 ~f:(fun (pc', _) -> compile_block blocks debug_data code pc' state');
-        Array.iter l2 ~f:(fun (pc', _) -> compile_block blocks debug_data code pc' state')
-    | Pushtrap _ | Raise _ | Return _ | Stop -> ())
+    | Switch (_, _) -> ()
+    | Pushtrap _ -> ()
+    | Raise _ | Return _ | Stop -> ())
 
 and compile infos pc state instrs =
   if debug_parser () then State.print state;
@@ -1630,20 +1629,67 @@ and compile infos pc state instrs =
         let x, _ = State.accu state in
         let args = State.stack_vars state in
         instrs, (Cond (x, (pc + 2, args), (pc + offset + 1, args)), loc), state
-    | SWITCH ->
+    | SWITCH -> (
         if debug_parser () then Format.printf "switch ...@.";
 
         let sz = getu code (pc + 1) in
         let x, _ = State.accu state in
         let args = State.stack_vars state in
         let l = sz land 0xFFFF in
+        let isize = sz land 0XFFFF in
+        let bsize = sz lsr 16 in
         let it =
           Array.init (sz land 0XFFFF) ~f:(fun i -> pc + 2 + gets code (pc + 2 + i), args)
         in
         let bt =
           Array.init (sz lsr 16) ~f:(fun i -> pc + 2 + gets code (pc + 2 + l + i), args)
         in
-        instrs, (Switch (x, it, bt), loc), state
+        Array.iter it ~f:(fun (pc', _) ->
+            compile_block infos.blocks infos.debug code pc' state);
+        Array.iter bt ~f:(fun (pc', _) ->
+            compile_block infos.blocks infos.debug code pc' state);
+        match isize, bsize with
+        | _, 0 -> instrs, (Switch (x, it), loc), state
+        | 0, _ ->
+            let x_tag = Var.fresh () in
+            let instrs =
+              (Let (x_tag, Prim (Extern "%direct_obj_tag", [ Pv x ])), loc) :: instrs
+            in
+            instrs, (Switch (x_tag, bt), loc), state
+        | _, _ ->
+            let isint_branch = pc + 1 in
+            let isblock_branch = pc + 2 in
+            let i_state = State.start_block isint_branch state in
+            let i_args = State.stack_vars i_state in
+            let b_state = State.start_block isblock_branch state in
+            let b_args = State.stack_vars b_state in
+            tagged_blocks := Addr.Set.add isint_branch !tagged_blocks;
+            tagged_blocks := Addr.Set.add isblock_branch !tagged_blocks;
+            let it = Array.map it ~f:(fun (x, _) -> x, i_args) in
+            let bt = Array.map bt ~f:(fun (x, _) -> x, b_args) in
+            let x_tag = Var.fresh () in
+            let () =
+              compiled_blocks :=
+                Addr.Map.add
+                  isint_branch
+                  (i_state, [], (Switch (x, it), loc))
+                  !compiled_blocks
+            in
+            let () =
+              let instrs =
+                [ Let (x_tag, Prim (Extern "%direct_obj_tag", [ Pv x ])), loc ]
+              in
+              compiled_blocks :=
+                Addr.Map.add
+                  isblock_branch
+                  (b_state, instrs, (Switch (x_tag, bt), loc))
+                  !compiled_blocks
+            in
+            let isint_var = Var.fresh () in
+            let instrs = (Let (isint_var, Prim (IsInt, [ Pv x ])), loc) :: instrs in
+            ( instrs
+            , (Cond (isint_var, (isint_branch, args), (isblock_branch, args)), loc)
+            , state ))
     | BOOLNOT ->
         let y, _ = State.accu state in
         let x, state = State.fresh_var state loc in