More precise tail-call analysis

Author: vouillon

compiler/lib-wasm/call_graph_analysis.ml

@@ -42,11 +42,35 @@ type t = { unambiguous_non_escaping : unit Var.Hashtbl.t }
 let direct_calls_only info f =
   Config.Flag.optcall () && Var.Hashtbl.mem info.unambiguous_non_escaping f
 
-let call_graph p info call_info =
+let callee_if_known info call_info exact f =
+  match get_approx info f with
+  | Top -> None
+  | Values { known; others } ->
+      if
+        exact
+        && (not others)
+        && Var.Set.for_all (fun f -> direct_calls_only call_info f) known
+      then Some (Var.Set.choose known)
+      else None
+
+let propagate nodes edges eligible =
+  let rec propagate n =
+    List.iter
+      ~f:(fun n' ->
+        if (not (Var.Hashtbl.mem nodes n')) && eligible n'
+        then (
+          Var.Hashtbl.add nodes n' ();
+          propagate n'))
+      (Var.Hashtbl.find_all edges n)
+  in
+  Var.Hashtbl.iter (fun n () -> propagate n) nodes
+
+let call_graph p info call_info eligible =
   let under_handler = Var.Hashtbl.create 16 in
   let callees = Var.Hashtbl.create 16 in
   let callers = Var.Hashtbl.create 16 in
   let has_tail_calls = Var.Hashtbl.create 16 in
+  let tail_callers = Var.Hashtbl.create 16 in
   let rec traverse name_opt pc visited nesting =
     if not (Addr.Set.mem pc visited)
     then (
@@ -92,8 +116,11 @@ let call_graph p info call_info =
                 match block.branch with
                 | Return x -> (
                     match last_instr block.body with
-                    | Some (Let (x', Apply _)) when Code.Var.equal x x' ->
-                        Var.Hashtbl.replace has_tail_calls f ()
+                    | Some (Let (x', Apply { f = g; exact; _ })) when Code.Var.equal x x'
+                      -> (
+                        match callee_if_known info call_info exact g with
+                        | None -> Var.Hashtbl.replace has_tail_calls f ()
+                        | Some g -> Var.Hashtbl.add tail_callers g f)
                     | _ -> ())
                 | _ -> ())
               pc
@@ -118,6 +145,7 @@ let call_graph p info call_info =
     p
     (fun name_opt _ (pc, _) _ () -> ignore (traverse name_opt pc Addr.Set.empty 0))
     ();
+  propagate has_tail_calls tail_callers eligible;
   under_handler, callers, callees, has_tail_calls
 
 let function_do_raise p pc =
@@ -134,20 +162,10 @@ let function_do_raise p pc =
     p.blocks
     false
 
-let propagate nodes edges eligible =
-  let rec propagate n =
-    List.iter
-      ~f:(fun n' ->
-        if (not (Var.Hashtbl.mem nodes n')) && eligible n'
-        then (
-          Var.Hashtbl.add nodes n' ();
-          propagate n'))
-      (Var.Hashtbl.find_all edges n)
-  in
-  Var.Hashtbl.iter (fun n () -> propagate n) nodes
-
 let raising_functions p info call_info eligible =
-  let under_handler, callers, callees, has_tail_calls = call_graph p info call_info in
+  let under_handler, callers, callees, has_tail_calls =
+    call_graph p info call_info eligible
+  in
   propagate under_handler callees (fun f ->
       eligible f && not (Var.Hashtbl.mem has_tail_calls f));
   let h = Var.Hashtbl.create 16 in