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

@@ -189,6 +189,46 @@ let is_int info x =
   | Pc (Int _) -> Y
   | Pc _ -> N
 
+let the_tag_of info x get =
+  match x with
+  | Pv x ->
+      get_approx
+        info
+        (fun x ->
+          match info.info_defs.(Var.idx x) with
+          | Expr (Block (j, _, _)) ->
+              if Var.ISet.mem info.info_possibly_mutable x then None else get j
+          | Expr (Constant (Tuple (j, _, _))) -> get j
+          | _ -> None)
+        None
+        (fun u v ->
+          match u, v with
+          | Some i, Some j when Poly.(i = j) -> u
+          | _ -> None)
+        x
+  | Pc (Tuple (j, _, _)) -> get j
+  | _ -> None
+
+let the_cont_of info x (a : cont array) =
+  (* The value of [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 *)
+  let get i = if i >= 0 && i < Array.length a then Some a.(i) else None in
+  get_approx
+    info
+    (fun x ->
+      match info.info_defs.(Var.idx x) with
+      | Expr (Prim (Extern "%direct_obj_tag", [ b ])) -> the_tag_of info b get
+      | Expr (Constant (Int j)) -> get (Int32.to_int j)
+      | _ -> None)
+    None
+    (fun u v ->
+      match u, v with
+      | Some i, Some j when Poly.(i = j) -> u
+      | _ -> None)
+    x
+
 let eval_instr info ((x, loc) as i) =
   match x with
   | Let (x, Prim (Extern ("caml_js_equals" | "caml_equal"), [ y; z ])) -> (
@@ -228,6 +268,13 @@ let eval_instr info ((x, loc) as i) =
           let c = Constant (Int b) in
           Flow.update_def info x c;
           [ Let (x, c), loc ])
+  | Let (x, Prim (Extern "%direct_obj_tag", [ y ])) -> (
+      match the_tag_of info y (fun x -> Some x) with
+      | Some tag ->
+          let c = Constant (Int (Int32.of_int tag)) in
+          Flow.update_def info x c;
+          [ Let (x, c), loc ]
+      | None -> [ i ])
   | Let (x, Prim (Extern "caml_sys_const_backend_type", [ _ ])) ->
       let jsoo = Code.Var.fresh () in
       [ Let (jsoo, Constant (String "js_of_ocaml")), noloc
@@ -271,34 +318,6 @@ let eval_instr info ((x, loc) as i) =
           ])
   | _ -> [ i ]
 
-type case_of =
-  | CConst of int
-  | CTag of int
-  | Unknown
-
-let the_case_of info x =
-  match x with
-  | Pv x ->
-      get_approx
-        info
-        (fun x ->
-          match info.info_defs.(Var.idx x) with
-          | Expr (Constant (Int i)) -> CConst (Int32.to_int i)
-          | Expr (Block (j, _, _)) ->
-              if Var.ISet.mem info.info_possibly_mutable x then Unknown else CTag j
-          | Expr (Constant (Tuple (j, _, _))) -> CTag j
-          | _ -> Unknown)
-        Unknown
-        (fun u v ->
-          match u, v with
-          | CTag i, CTag j when i = j -> u
-          | CConst i, CConst j when i = j -> u
-          | _ -> Unknown)
-        x
-  | Pc (Int i) -> CConst (Int32.to_int i)
-  | Pc (Tuple (j, _, _)) -> CTag j
-  | _ -> Unknown
-
 type cond_of =
   | Zero
   | Non_zero
@@ -341,15 +360,10 @@ let eval_branch info (l, loc) =
           | Zero -> Branch ffalse
           | Non_zero -> Branch ftrue
           | Unknown -> b)
-    | Switch (x, const, tags) 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)
+    | Switch (x, a) as b -> (
+        match the_cont_of info x a with
+        | Some cont -> Branch cont
+        | None -> b)
     | _ as b -> b
   in
   l, loc
@@ -380,15 +394,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

@@ -911,6 +911,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));
@@ -1440,15 +1441,10 @@ and compile_block_no_loop st queue (pc : Addr.t) ~fall_through scope_stack =
   let seq, queue = translate_instrs st.ctx queue block.body block.branch in
   let nbbranch =
     match fst block.branch with
-    | Switch (_, a, b) ->
+    | Switch (_, a) ->
         (* Build an artifical dtree with the correct layout so that
            [Dtree.nbbranch dtree pc] is correct *)
-        let dtree =
-          match a, b with
-          | [||], _ -> DTree.build_switch b
-          | _, [||] -> DTree.build_switch a
-          | _ -> DTree.If (IsTrue, DTree.build_switch a, DTree.build_switch b)
-        in
+        let dtree = DTree.build_switch a in
         fun pc -> DTree.nbbranch dtree pc
     | Cond (_, a, b) ->
         let dtree = DTree.build_if a b in
@@ -1583,7 +1579,7 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
      | 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
@@ -1642,20 +1638,7 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
             (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
-            "Tag"
-            st
-            scope_stack
-            ~fall_through
-            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
@@ -1668,39 +1651,6 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
             (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
-            "Int"
-            st
-            scope_stack
-            ~fall_through
-            loc
-            (var x)
-            (DTree.build_switch a1)
-        in
-        let never2, b2 =
-          compile_decision_tree
-            "Tag"
-            st
-            scope_stack
-            ~fall_through
-            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