Another liveness hack

asmcomp/liveness.ml

@@ -22,12 +22,41 @@ type liveness_env =
   { at_exit : (int * Reg.Set.t) list;
     at_raise : Reg.Set.t;
     last_regular_trywith_handler : Reg.Set.t;
+    free_conts_for_handlers : Numbers.Int.Set.t Numbers.Int.Map.t;
   }
 
-let initial_env =
+let initial_env fundecl =
   { at_exit = [];
     at_raise = Reg.Set.empty;
     last_regular_trywith_handler = Reg.Set.empty;
+    free_conts_for_handlers = Mach.free_conts_for_handlers fundecl;
+  }
+
+let same_env env1 env2 =
+  List.for_all2 (fun (nfail1, live_regs1) (nfail2, live_regs2) ->
+      nfail1 = nfail2 && Reg.Set.equal live_regs1 live_regs2)
+    env1.at_exit
+    env2.at_exit
+  && Reg.Set.equal env1.at_raise env2.at_raise
+  && Reg.Set.equal env1.last_regular_trywith_handler
+       env2.last_regular_trywith_handler
+
+type cache_entry =
+  { restricted_env : liveness_env;  (* last used environment,
+                                       restricted to the live conts *)
+    at_join : Reg.Set.t;            (* last used set at join *)
+    before_handler : Reg.Set.t;    (* last computed result *)
+}
+
+let fixpoint_cache : cache_entry Numbers.Int.Map.t ref =
+ ref Numbers.Int.Map.empty
+
+let reset_cache () = fixpoint_cache := Numbers.Int.Map.empty
+
+let restrict_env env conts =
+  { env with at_exit =
+               List.filter (fun (n, _) -> Numbers.Int.Set.mem n conts)
+                 env.at_exit;
   }
 
 let find_live_at_exit env k =
@@ -113,7 +142,28 @@ let rec live env i finally =
       let aux env (nfail, ts, handler) (nfail', before_handler) =
         assert(nfail = nfail');
         let env = env_from_trap_stack env ts in
-        let before_handler' = live env handler at_join in
+        let free_conts = Numbers.Int.Map.find nfail env.free_conts_for_handlers in
+        let before_handler', restricted_env, do_update =
+          match Numbers.Int.Map.find nfail !fixpoint_cache with
+          | exception Not_found ->
+            let restricted_env = restrict_env env free_conts in
+            live env handler at_join, restricted_env, true
+          | cache ->
+            let restricted_env = restrict_env env free_conts in
+            if same_env restricted_env cache.restricted_env
+            && Reg.Set.equal at_join cache.at_join
+            then cache.before_handler, cache.restricted_env, false
+            else live env handler at_join, restricted_env, true
+        in
+        if do_update then begin
+          let cache_entry =
+            { restricted_env;
+              at_join;
+              before_handler = before_handler';
+            }
+          in
+          fixpoint_cache := Numbers.Int.Map.add nfail cache_entry !fixpoint_cache
+        end;
         nfail, Reg.Set.union before_handler before_handler'
       in
       let aux_equal (nfail, before_handler) (nfail', before_handler') =
@@ -168,7 +218,8 @@ let rec live env i finally =
       Reg.add_set_array env.at_raise arg
 
 let fundecl f =
-  let initially_live = live initial_env f.fun_body Reg.Set.empty in
+  reset_cache ();
+  let initially_live = live (initial_env f) f.fun_body Reg.Set.empty in
   (* Sanity check: only function parameters (and the Spacetime node hole
      register, if profiling) can be live at entrypoint *)
   let wrong_live = Reg.Set.diff initially_live (Reg.set_of_array f.fun_args) in

asmcomp/mach.ml

@@ -210,3 +210,53 @@ let operation_can_raise op =
   | Iintop (Icheckbound _) | Iintop_imm (Icheckbound _, _)
   | Ialloc _ -> true
   | _ -> false
+
+let free_conts_for_handlers fundecl =
+  let module S = Numbers.Int.Set in
+  let module M = Numbers.Int.Map in
+  let acc = ref M.empty in
+  let rec free_conts i =
+    match i.desc with
+    | Iend -> S.empty
+    | desc ->
+      let next_conts = free_conts i.next in
+      match desc with
+      | Iend -> assert false
+      | Iop _ -> next_conts
+      | Ireturn _ -> next_conts
+      | Iifthenelse (_, then_, else_) ->
+        S.union next_conts (S.union (free_conts then_) (free_conts else_))
+      | Iswitch (_, cases) ->
+        Array.fold_left (fun conts instr -> S.union conts (free_conts instr))
+          next_conts cases
+      | Icatch (_rec_flag, handlers, body) ->
+        let conts = S.union next_conts (free_conts body) in
+        let conts =
+          List.fold_left (fun conts (nfail, ts, i) ->
+            let rec add_exn_conts conts = function
+              | Uncaught -> conts
+              | Generic_trap ts -> add_exn_conts conts ts
+              | Specific_trap (nfail, ts) -> add_exn_conts (S.add nfail conts) ts
+            in
+            let free = add_exn_conts (free_conts i) ts in
+            acc := M.add nfail free !acc;
+            S.union conts free)
+            conts handlers
+        in
+        List.fold_left (fun conts (nfail, _ts, _i) ->
+          S.remove nfail conts)
+          conts handlers
+      | Iexit (nfail, _) -> S.add nfail next_conts
+      | Itrywith (body, kind, (_ts, handler)) ->
+        let conts =
+          S.union next_conts (S.union (free_conts body) (free_conts handler))
+        in
+        begin match kind with
+        | Regular -> conts
+        | Delayed nfail -> S.remove nfail conts
+        end
+      | Iraise _ -> next_conts
+  in
+  let free = free_conts fundecl.fun_body in
+  assert(S.is_empty free);
+  !acc

asmcomp/mach.mli

@@ -147,3 +147,5 @@ val instr_iter: (instruction -> unit) -> instruction -> unit
 val spacetime_node_hole_pointer_is_live_before : instruction -> bool
 
 val operation_can_raise : operation -> bool
+
+val free_conts_for_handlers : fundecl -> Numbers.Int.Set.t Numbers.Int.Map.t