Compiler: use shapes to compute arity

This reverts commit 9c0692f8af4d6fafdf0fc428f820cc9753897cc4.

Author: hhugo

compiler/lib/driver.ml

@@ -65,7 +65,10 @@ let inline profile p =
 
 let specialize_1 (p, info) =
   if debug () then Format.eprintf "Specialize...@.";
-  Specialize.f ~function_arity:(fun f -> Specialize.function_arity info f) p
+  let return_values = Code.Var.Map.empty in
+  Specialize.f
+    ~function_arity:(fun f -> Specialize.function_arity ~return_values info f)
+    p
 
 let specialize_js (p, info) =
   if debug () then Format.eprintf "Specialize js...@.";

compiler/lib/specialize.ml

@@ -19,41 +19,17 @@
  *)
 open! Stdlib
 open Code
-open Flow
 
 let times = Debug.find "times"
 
 let stats = Debug.find "stats"
 
 let debug_stats = Debug.find "stats-debug"
 
-let function_arity info x =
-  let rec arity info x acc =
-    get_approx
-      info
-      (fun x ->
-        match Flow.Info.def info x with
-        | Some (Closure (l, _, _)) -> Some (List.length l)
-        | Some (Special (Alias_prim prim)) -> (
-            try Some (Primitive.arity prim) with Not_found -> None)
-        | Some (Apply { f; args; _ }) -> (
-            if List.mem ~eq:Var.equal f 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 function_arity ~return_values info x =
+  match Flow.the_shape_of ~return_values ~pure:Pure_fun.empty info x with
+  | Top | Block _ -> None
+  | Function { arity; _ } -> Some arity
 
 let add_event loc instrs =
   match loc with

compiler/lib/specialize.mli

@@ -18,7 +18,8 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-val function_arity : Flow.Info.t -> Code.Var.t -> int option
+val function_arity :
+  return_values:Code.Var.Set.t Code.Var.Map.t -> Flow.Info.t -> Code.Var.t -> int option
 
 val f : function_arity:(Code.Var.t -> int option) -> Code.program -> Code.program