Compiler: remove last argument of Pushtrap (#1575)

Author: hhugo

compiler/lib/code.ml

@@ -360,7 +360,7 @@ type last =
   | Branch of cont
   | Cond of Var.t * cont * cont
   | Switch of Var.t * cont array
-  | Pushtrap of cont * Var.t * cont * Addr.Set.t
+  | Pushtrap of cont * Var.t * cont
   | Poptrap of cont
 
 type block =
@@ -514,17 +514,8 @@ module Print = struct
         Format.fprintf f "switch %a {" Var.print x;
         Array.iteri a1 ~f:(fun i c -> Format.fprintf f "int %d -> %a; " i cont c);
         Format.fprintf f "}"
-    | Pushtrap (cont1, x, cont2, pcs) ->
-        Format.fprintf
-          f
-          "pushtrap %a handler %a => %a continuation %s"
-          cont
-          cont1
-          Var.print
-          x
-          cont
-          cont2
-          (String.concat ~sep:", " (List.map (Addr.Set.elements pcs) ~f:string_of_int))
+    | Pushtrap (cont1, x, cont2) ->
+        Format.fprintf f "pushtrap %a handler %a => %a" cont cont1 Var.print x cont cont2
     | Poptrap c -> Format.fprintf f "poptrap %a" cont c
 
   type xinstr =
@@ -600,7 +591,7 @@ let fold_children blocks pc f accu =
   match fst block.branch with
   | Return _ | Raise _ | Stop -> accu
   | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
-  | Pushtrap ((pc', _), _, (pc_h, _), _) ->
+  | Pushtrap ((pc', _), _, (pc_h, _)) ->
       let accu = f pc' accu in
       let accu = f pc_h accu in
       accu
@@ -635,6 +626,39 @@ let rec traverse' { fold } f pc visited blocks acc =
 
 let traverse fold f pc blocks acc = snd (traverse' fold f pc Addr.Set.empty blocks acc)
 
+let poptraps blocks pc =
+  let rec loop blocks pc visited depth acc =
+    if Addr.Set.mem pc visited
+    then acc, visited
+    else
+      let visited = Addr.Set.add pc visited in
+      let block = Addr.Map.find pc blocks in
+      match fst block.branch with
+      | Return _ | Raise _ | Stop -> acc, visited
+      | Branch (pc', _) -> loop blocks pc' visited depth acc
+      | Poptrap (pc', _) ->
+          if depth = 0
+          then Addr.Set.add pc' acc, visited
+          else loop blocks pc' visited (depth - 1) acc
+      | Pushtrap ((pc', _), _, (pc_h, _)) ->
+          let acc, visited = loop blocks pc' visited (depth + 1) acc in
+          let acc, visited = loop blocks pc_h visited depth acc in
+          acc, visited
+      | Cond (_, (pc1, _), (pc2, _)) ->
+          let acc, visited = loop blocks pc1 visited depth acc in
+          let acc, visited = loop blocks pc2 visited depth acc in
+          acc, visited
+      | Switch (_, a) ->
+          let acc, visited =
+            Array.fold_right
+              ~init:(acc, visited)
+              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
+              a
+          in
+          acc, visited
+  in
+  loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
+
 let rec preorder_traverse' { fold } f pc visited blocks acc =
   if not (Addr.Set.mem pc visited)
   then
@@ -737,7 +761,7 @@ let invariant { blocks; start; _ } =
           check_cont cont1;
           check_cont cont2
       | Switch (_x, a1) -> Array.iteri a1 ~f:(fun _ cont -> check_cont cont)
-      | Pushtrap (cont1, _x, cont2, _pcs) ->
+      | Pushtrap (cont1, _x, cont2) ->
           check_cont cont1;
           check_cont cont2
       | Poptrap cont -> check_cont cont

compiler/lib/code.mli

@@ -207,7 +207,7 @@ type last =
   | Branch of cont
   | Cond of Var.t * cont * cont
   | Switch of Var.t * cont array
-  | Pushtrap of cont * Var.t * cont * Addr.Set.t
+  | Pushtrap of cont * Var.t * cont
   | Poptrap of cont
 
 type block =
@@ -267,6 +267,8 @@ val fold_closures_innermost_first :
 
 val fold_children : 'c fold_blocs
 
+val poptraps : block Addr.Map.t -> Addr.t -> Addr.Set.t
+
 val traverse :
   fold_blocs_poly -> (Addr.t -> 'c -> 'c) -> Addr.t -> block Addr.Map.t -> 'c -> 'c
 

compiler/lib/deadcode.ml

@@ -115,7 +115,7 @@ and mark_reachable st pc =
     | Switch (x, a1) ->
         mark_var st x;
         Array.iter a1 ~f:(fun cont -> mark_cont_reachable st cont)
-    | Pushtrap (cont1, _, cont2, _) ->
+    | Pushtrap (cont1, _, cont2) ->
         mark_cont_reachable st cont1;
         mark_cont_reachable st cont2)
 
@@ -152,12 +152,8 @@ let filter_live_last blocks st (l, loc) =
         Cond (x, filter_cont blocks st cont1, filter_cont blocks st cont2)
     | 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
-          , x
-          , filter_cont blocks st cont2
-          , Addr.Set.inter pcs st.reachable_blocks )
+    | Pushtrap (cont1, x, cont2) ->
+        Pushtrap (filter_cont blocks st cont1, x, filter_cont blocks st cont2)
     | Poptrap cont -> Poptrap (filter_cont blocks st cont)
   in
   l, loc
@@ -213,7 +209,7 @@ let f ({ blocks; _ } as p : Code.program) =
           add_cont_dep blocks defs cont1;
           add_cont_dep blocks defs cont2
       | Switch (_, a1) -> Array.iter a1 ~f:(fun cont -> add_cont_dep blocks defs cont)
-      | Pushtrap (cont, _, cont_h, _) ->
+      | Pushtrap (cont, _, cont_h) ->
           add_cont_dep blocks defs cont_h;
           add_cont_dep blocks defs cont
       | Poptrap cont -> add_cont_dep blocks defs cont)

compiler/lib/effects.ml

@@ -191,7 +191,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
               List.iter ~f:mark_needed englobing_exn_handlers;
               mark_continuation dst x
           | _ -> ())
-      | Pushtrap (_, x, (handler_pc, _), _) -> mark_continuation handler_pc x
+      | Pushtrap (_, x, (handler_pc, _)) -> mark_continuation handler_pc x
       | Poptrap _ | Raise _ -> (
           match englobing_exn_handlers with
           | handler_pc :: _ -> Hashtbl.add matching_exn_handler pc handler_pc
@@ -203,7 +203,7 @@ let compute_needed_transformations ~cfg ~idom ~cps_needed ~blocks ~start =
         (fun pc visited ->
           let englobing_exn_handlers =
             match block.branch with
-            | Pushtrap (_, _, (handler_pc, _), _), _ when pc <> handler_pc ->
+            | Pushtrap (_, _, (handler_pc, _)), _ when pc <> handler_pc ->
                 handler_pc :: englobing_exn_handlers
             | Poptrap _, _ -> List.tl englobing_exn_handlers
             | _ -> englobing_exn_handlers
@@ -423,7 +423,7 @@ let cps_last ~st ~alloc_jump_closures pc ((last, last_loc) : last * loc) ~k :
          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), last_loc)
-  | Pushtrap (body_cont, exn, ((handler_pc, _) as handler_cont), _) -> (
+  | 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
       | false -> alloc_jump_closures, (last, last_loc)
@@ -910,8 +910,7 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
                | Branch cont -> Branch (resolve cont)
                | Cond (x, cont1, cont2) -> Cond (x, resolve cont1, resolve cont2)
                | Switch (x, a1) -> Switch (x, Array.map ~f:resolve a1)
-               | Pushtrap (cont1, x, cont2, s) ->
-                   Pushtrap (resolve cont1, x, resolve cont2, s)
+               | Pushtrap (cont1, x, cont2) -> Pushtrap (resolve cont1, x, resolve cont2)
                | Poptrap cont -> Poptrap (resolve cont)
                | Return _ | Raise _ | Stop -> branch
              in

compiler/lib/eval.ml

@@ -407,8 +407,7 @@ let drop_exception_handler blocks =
   Addr.Map.fold
     (fun pc _ blocks ->
       match Addr.Map.find pc blocks with
-      | { branch = Pushtrap (((addr, _) as cont1), _x, _cont2, addrset), loc; _ } as b
-        -> (
+      | { branch = Pushtrap (((addr, _) as cont1), _x, _cont2), loc; _ } as b -> (
           try
             let visited = do_not_raise addr Addr.Set.empty blocks in
             let b = { b with branch = Branch cont1, loc } in
@@ -419,9 +418,7 @@ let drop_exception_handler blocks =
                   let b = Addr.Map.find pc2 blocks in
                   let branch =
                     match b.branch with
-                    | Poptrap ((addr, _) as cont), loc ->
-                        assert (Addr.Set.mem addr addrset);
-                        Branch cont, loc
+                    | Poptrap cont, loc -> Branch cont, loc
                     | x -> x
                   in
                   let b = { b with branch } in

compiler/lib/flow.ml

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

compiler/lib/freevars.ml

@@ -68,7 +68,7 @@ let iter_last_free_var f l =
   | Switch (x, a1) ->
       f x;
       Array.iter a1 ~f:(fun c -> iter_cont_free_vars f c)
-  | Pushtrap (cont1, _, cont2, _) ->
+  | Pushtrap (cont1, _, cont2) ->
       iter_cont_free_vars f cont1;
       iter_cont_free_vars f cont2
 
@@ -84,7 +84,7 @@ let iter_instr_bound_vars f i =
 let iter_last_bound_vars f l =
   match l with
   | Return _ | Raise _ | Stop | Branch _ | Cond _ | Switch _ | Poptrap _ -> ()
-  | Pushtrap (_, x, _, _) -> f x
+  | Pushtrap (_, x, _) -> f x
 
 let iter_block_bound_vars f block =
   List.iter ~f block.params;

compiler/lib/generate.ml

@@ -1704,7 +1704,7 @@ and compile_conditional st queue ~fall_through last scope_stack : _ * _ =
         in
         true, flush_all queue [ J.Return_statement e_opt, loc ]
     | Branch cont -> compile_branch st queue cont scope_stack ~fall_through
-    | Pushtrap (c1, x, e1, _) ->
+    | Pushtrap (c1, x, e1) ->
         let never_body, body = compile_branch st [] c1 scope_stack ~fall_through in
         if debug () then Format.eprintf "@,}@]@,@[<hv 2>catch {@;";
         let never_handler, handler = compile_branch st [] e1 scope_stack ~fall_through in

compiler/lib/global_deadcode.ml

@@ -158,7 +158,7 @@ let usages prog (global_info : Global_flow.info) : usage_kind Var.Map.t Var.Tbl.
           add_cont_deps cont1;
           add_cont_deps cont2
       | Switch (_, a) -> Array.iter ~f:add_cont_deps a
-      | Pushtrap (cont, _, cont_h, _) ->
+      | Pushtrap (cont, _, cont_h) ->
           add_cont_deps cont;
           add_cont_deps cont_h
       | Poptrap cont -> add_cont_deps cont)
@@ -380,7 +380,7 @@ let zero prog sentinal live_table =
       | Branch _, _
       | Cond (_, _, _), _
       | Switch (_, _), _
-      | Pushtrap (_, _, _, _), _
+      | Pushtrap (_, _, _), _
       | Poptrap _, _ -> block.branch
     in
     { block with body; branch }

compiler/lib/global_flow.ml

@@ -275,7 +275,7 @@ let program_deps st { blocks; _ } =
                       block.body)
                   h
             | Expr _ | Phi _ -> ())
-      | Pushtrap (cont, x, cont_h, _) ->
+      | Pushtrap (cont, x, cont_h) ->
           add_var st x;
           st.defs.(Var.idx x) <- Phi { known = Var.Set.empty; others = true };
           cont_deps blocks st cont_h;

compiler/lib/inline.ml

@@ -107,7 +107,8 @@ let fold_children blocks pc f accu =
   match fst block.branch with
   | Return _ | Raise _ | Stop -> accu
   | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
-  | Pushtrap (_, _, (pc1, _), pcs) -> f pc1 (Addr.Set.fold f pcs accu)
+  | Pushtrap ((try_body, _), _, (pc1, _)) ->
+      f pc1 (Addr.Set.fold f (Code.poptraps blocks try_body) accu)
   | Cond (_, (pc1, _), (pc2, _)) ->
       let accu = f pc1 accu in
       let accu = f pc2 accu in