More exact calls

Author: vouillon

compiler/lib/driver.ml

@@ -49,7 +49,7 @@ let inline p =
 
 let specialize_1 (p, info) =
   if debug () then Format.eprintf "Specialize...@.";
-  Specialize.f info p
+  Specialize.f ~function_arity:(fun f -> Specialize.function_arity info f) p
 
 let specialize_js (p, info) =
   if debug () then Format.eprintf "Specialize js...@.";
@@ -91,6 +91,13 @@ let effects p =
     p |> Deadcode.f +> Effects.f +> map_fst Lambda_lifting.f)
   else p, (Code.Var.Set.empty : Effects.cps_calls)
 
+let exact_calls p =
+  if not (Config.Flag.effects ())
+  then
+    let info = Global_flow.f p in
+    Specialize.f ~function_arity:(fun f -> Global_flow.function_arity info f) p
+  else p
+
 let print p =
   if debug () then Code.Print.program (fun _ _ -> "") p;
   p
@@ -565,7 +572,11 @@ let full
     p =
   let exported_runtime = not standalone in
   let opt =
-    specialize_js_once +> profile +> effects +> map_fst (Generate_closure.f +> deadcode')
+    specialize_js_once
+    +> profile
+    +> exact_calls
+    +> effects
+    +> map_fst (Generate_closure.f +> deadcode')
   in
   let emit =
     generate d ~exported_runtime ~wrap_with_fun ~warn_on_unhandled_effect:standalone

compiler/lib/global_flow.ml

@@ -621,3 +621,25 @@ let exact_call info f n =
           | Expr (Block _) -> true
           | Expr _ | Phi _ -> assert false)
         known
+
+let function_arity info f =
+  match Var.Tbl.get info.info_approximation f with
+  | Top | Values { others = true; _ } -> None
+  | Values { known; others = false } -> (
+      match
+        Var.Set.fold
+          (fun g acc ->
+            match info.info_defs.(Var.idx g) with
+            | Expr (Closure (params, _)) -> (
+                let n = List.length params in
+                match acc with
+                | None -> Some (Some n)
+                | Some (Some n') when n <> n' -> Some None
+                | Some _ -> acc)
+            | Expr (Block _) -> acc
+            | Expr _ | Phi _ -> assert false)
+          known
+          None
+      with
+      | Some v -> v
+      | None -> None)

compiler/lib/global_flow.mli

@@ -40,3 +40,5 @@ type info =
 val f : Code.program -> info
 
 val exact_call : info -> Var.t -> int -> bool
+
+val function_arity : info -> Var.t -> int option

compiler/lib/specialize.ml

@@ -21,36 +21,39 @@ open! Stdlib
 open Code
 open Flow
 
-let rec function_cardinality info x acc =
-  get_approx
-    info
-    (fun x ->
-      match info.info_defs.(Var.idx x) with
-      | Expr (Closure (l, _)) -> Some (List.length l)
-      | Expr (Prim (Extern "%closure", [ Pc (String prim) ])) -> (
-          try Some (Primitive.arity prim) with Not_found -> None)
-      | Expr (Apply { f; args; _ }) -> (
-          if List.mem f ~set:acc
-          then None
-          else
-            match function_cardinality info f (f :: acc) with
-            | Some n ->
-                let diff = n - List.length args in
-                if diff > 0 then Some diff else None
-            | None -> None)
-      | _ -> None)
-    None
-    (fun u v ->
-      match u, v with
-      | Some n, Some m when n = m -> u
-      | _ -> None)
-    x
+let function_arity info x =
+  let rec arity info x acc =
+    get_approx
+      info
+      (fun x ->
+        match info.info_defs.(Var.idx x) with
+        | Expr (Closure (l, _)) -> Some (List.length l)
+        | Expr (Prim (Extern "%closure", [ Pc (String prim) ])) -> (
+            try Some (Primitive.arity prim) with Not_found -> None)
+        | Expr (Apply { f; args; _ }) -> (
+            if List.mem f ~set:acc
+            then None
+            else
+              match arity info f (f :: acc) with
+              | Some n ->
+                  let diff = n - List.length args in
+                  if diff > 0 then Some diff else None
+              | None -> None)
+        | _ -> None)
+      None
+      (fun u v ->
+        match u, v with
+        | Some n, Some m when n = m -> u
+        | _ -> None)
+      x
+  in
+  arity info x []
 
-let specialize_instr info (acc, free_pc, extra) i =
+let specialize_instr function_arity (acc, free_pc, extra) i =
   match i with
-  | Let (x, Apply { f; args; _ }) when Config.Flag.optcall () -> (
+  | Let (x, Apply { f; args; exact }) when (not exact) && Config.Flag.optcall () -> (
       let n' = List.length args in
-      match function_cardinality info f [] with
+      match function_arity f with
       | None -> i :: acc, free_pc, extra
       | Some n when n = n' ->
           Let (x, Apply { f; args; exact = true }) :: acc, free_pc, extra
@@ -80,13 +83,13 @@ let specialize_instr info (acc, free_pc, extra) i =
       | _ -> i :: acc, free_pc, extra)
   | _ -> i :: acc, free_pc, extra
 
-let specialize_instrs info p =
+let specialize_instrs ~function_arity p =
   let blocks, free_pc =
     Addr.Map.fold
       (fun pc block (blocks, free_pc) ->
         let body, free_pc, extra =
           List.fold_right block.body ~init:([], free_pc, []) ~f:(fun i acc ->
-              specialize_instr info acc i)
+              specialize_instr function_arity acc i)
         in
         let blocks =
           List.fold_left extra ~init:blocks ~f:(fun blocks (pc, b) ->
@@ -98,4 +101,4 @@ let specialize_instrs info p =
   in
   { p with blocks; free_pc }
 
-let f info p = specialize_instrs info p
+let f = specialize_instrs

compiler/lib/specialize.mli

@@ -18,4 +18,6 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-val f : Flow.info -> Code.program -> Code.program
+val function_arity : Flow.info -> Code.Var.t -> int option
+
+val f : function_arity:(Code.Var.t -> int option) -> Code.program -> Code.program