Compiler: remove last argument of Pushtrap (#1575)

Author: hhugo

compiler/lib/code.ml

@@ -365,7 +365,7 @@ type last =
   | Branch of cont
   | Cond of Var.t * cont * cont
   | Switch of Var.t * cont array * cont array
-  | Pushtrap of cont * Var.t * cont * Addr.Set.t
+  | Pushtrap of cont * Var.t * cont
   | Poptrap of cont
 
 type block =
@@ -528,17 +528,8 @@ module Print = struct
         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
-          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 =
@@ -609,12 +600,51 @@ let is_empty p =
       | _ -> false)
   | _ -> false
 
+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 (_, a1, a2) ->
+          let acc, visited =
+            Array.fold_right
+              ~init:(acc, visited)
+              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
+              a1
+          in
+          let acc, visited =
+            Array.fold_right
+              ~init:(acc, visited)
+              ~f:(fun (pc, _) (acc, visited) -> loop blocks pc visited depth acc)
+              a2
+          in
+          acc, visited
+  in
+  loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
+
 let fold_children blocks pc f accu =
   let block = Addr.Map.find pc blocks in
   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
@@ -632,8 +662,8 @@ let fold_children_skip_try_body blocks pc f accu =
   match fst block.branch with
   | Return _ | Raise _ | Stop -> accu
   | Branch (pc', _) | Poptrap (pc', _) -> f pc' accu
-  | Pushtrap (_, _, (pc_h, _), pcs) ->
-      let accu = Addr.Set.fold f pcs accu in
+  | Pushtrap ((pc', _), _, (pc_h, _)) ->
+      let accu = Addr.Set.fold f (poptraps blocks pc') accu in
       let accu = f pc_h accu in
       accu
   | Cond (_, (pc1, _), (pc2, _)) ->
@@ -791,7 +821,7 @@ let invariant { blocks; start; _ } =
       | Switch (_x, a1, a2) ->
           Array.iteri a1 ~f:(fun _ cont -> check_cont cont);
           Array.iteri a2 ~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

@@ -212,7 +212,7 @@ type last =
   | Branch of cont
   | Cond of Var.t * cont * cont
   | Switch of Var.t * cont array * cont array
-  | Pushtrap of cont * Var.t * cont * Addr.Set.t
+  | Pushtrap of cont * Var.t * cont
   | Poptrap of cont
 
 type block =
@@ -276,6 +276,8 @@ val fold_children : 'c fold_blocs
 
 val fold_children_skip_try_body : '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

@@ -102,7 +102,7 @@ and mark_reachable st pc =
         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)
-    | Pushtrap (cont1, _, cont2, _) ->
+    | Pushtrap (cont1, _, cont2) ->
         mark_cont_reachable st cont1;
         mark_cont_reachable st cont2)
 
@@ -142,12 +142,8 @@ let filter_live_last blocks st (l, loc) =
           ( x
           , Array.map a1 ~f:(fun cont -> filter_cont blocks st cont)
           , Array.map a2 ~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
@@ -208,7 +204,7 @@ let f ({ blocks; _ } as p : Code.program) =
       | 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)
-      | 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
@@ -438,7 +438,7 @@ let cps_last ~st ~alloc_jump_closures pc ((last, last_loc) : last * loc) ~k :
       ( alloc_jump_closures
       , ( Switch (x, Array.map c1 ~f:cps_jump_cont, Array.map c2 ~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)
@@ -931,8 +931,7 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
                | 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)
-               | 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

@@ -444,8 +444,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
@@ -456,9 +455,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; _ } =
       | 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)
-      | 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

@@ -69,7 +69,7 @@ let iter_last_free_var f l =
       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)
-  | Pushtrap (cont1, _, cont2, _) ->
+  | Pushtrap (cont1, _, cont2) ->
       iter_cont_free_vars f cont1;
       iter_cont_free_vars f cont2
 
@@ -85,7 +85,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

@@ -758,7 +758,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 ((pc1, _), _, (pc2, _), _) ->
+  | Pushtrap ((pc1, _), _, (pc2, _)) ->
       let accu = f pc1 accu in
       let accu = f pc2 accu in
       accu
@@ -805,7 +805,7 @@ let build_graph ctx pc =
       List.iter pc_succs ~f:(fun pc' ->
           let pushtrap =
             match fst b.branch with
-            | Pushtrap ((pc1, _), _, (pc2, _), _remove) ->
+            | Pushtrap ((pc1, _), _, (pc2, _)) ->
                 if pc' = pc1
                 then (
                   Hashtbl.add poptrap pc Addr.Set.empty;
@@ -1866,7 +1866,7 @@ and compile_conditional st queue last loop_stack backs frontier interm =
         in
         true, flush_all queue [ J.Return_statement e_opt, loc ]
     | Branch cont -> compile_branch st queue cont loop_stack backs frontier interm
-    | Pushtrap (c1, x, e1, _) ->
+    | Pushtrap (c1, x, e1) ->
         let never_body, body = compile_branch st [] c1 loop_stack backs frontier interm in
         if debug () then Format.eprintf "@,}@]@,@[<hv 2>catch {@;";
         let never_handler, handler =

compiler/lib/global_flow.ml

@@ -273,7 +273,7 @@ let program_deps st { blocks; _ } =
                     | _ -> ())
                   block.body)
               h
-      | 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