Compiler: Specialze.f use shapes to compute arity

Author: hhugo

compiler/lib/driver.ml

@@ -67,7 +67,7 @@ let specialize_1 (p, info) =
   if debug () then Format.eprintf "Specialize...@.";
   let return_values = Code.Var.Map.empty in
   Specialize.f
-    ~function_arity:(fun f -> Specialize.function_arity ~return_values info f)
+    ~shape:(fun f -> Flow.the_shape_of ~return_values ~pure:Pure_fun.empty info f)
     ~update_def:(fun x expr -> Flow.Info.update_def info x expr)
     p
 
@@ -179,7 +179,10 @@ let effects_and_exact_calls
   | `Disabled | `Jspi ->
       let p =
         Specialize.f
-          ~function_arity:(fun f -> Global_flow.function_arity info f)
+          ~shape:(fun f ->
+            match Global_flow.function_arity info f with
+            | None -> Shape.Top
+            | Some arity -> Shape.Function { arity; pure = false; res = Top })
           ~update_def:(fun x expr -> Global_flow.update_def info x expr)
           p
       in

compiler/lib/specialize.ml

@@ -26,11 +26,6 @@ let stats = Debug.find "stats"
 
 let debug_stats = Debug.find "stats-debug"
 
-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
   | Some loc -> Event loc :: instrs
@@ -40,83 +35,91 @@ let unknown_apply = function
   | Let (_, Apply { f = _; args = _; exact = false }) -> true
   | _ -> false
 
-let specialize_apply opt_count function_arity update_def ((acc, free_pc, extra), loc) i =
-  match i with
-  | Let (x, Apply { f; args; exact = false }) -> (
-      let n' = List.length args in
-      match function_arity f with
-      | None -> i :: acc, free_pc, extra
-      | Some n when n = n' ->
-          incr opt_count;
-          Let (x, Apply { f; args; exact = true }) :: acc, free_pc, extra
-      | Some n when n < n' ->
+let specialize_apply opt_count shape update_def =
+  let rec loop x f args shape loc (acc, free_pc, extra) =
+    match (shape : Shape.t) with
+    | Top | Block _ -> Let (x, Apply { f; args; exact = false }) :: acc, free_pc, extra
+    | Function { arity; res; _ } ->
+        let nargs = List.length args in
+        if arity = nargs
+        then (
           incr opt_count;
-          let v = Code.Var.fresh () in
-          let args, rest = List.take n args in
-          ( (* Reversed *)
-            Let (x, Apply { f = v; args = rest; exact = false })
-            :: add_event loc (Let (v, Apply { f; args; exact = true }) :: acc)
-          , free_pc
-          , extra )
-      | Some n when n > n' ->
+          let expr = Apply { f; args; exact = true } in
+          update_def x expr;
+          Let (x, expr) :: acc, free_pc, extra)
+        else if arity > nargs
+        then (
+          (* under application *)
           incr opt_count;
-          let missing = Array.init (n - n') ~f:(fun _ -> Code.Var.fresh ()) in
+          let missing = Array.init (arity - nargs) ~f:(fun _ -> Code.Var.fresh ()) in
           let missing = Array.to_list missing in
           let block =
-            let params' = Array.init (n - n') ~f:(fun _ -> Code.Var.fresh ()) in
-            let params' = Array.to_list params' in
+            let params' = List.map missing ~f:Code.Var.fork in
             let return' = Code.Var.fresh () in
+            let args = args @ params' in
+            assert (List.length args = arity);
             { params = params'
-            ; body =
-                add_event
-                  loc
-                  [ Let (return', Apply { f; args = args @ params'; exact = true }) ]
+            ; body = add_event loc [ Let (return', Apply { f; args; exact = true }) ]
             ; branch = Return return'
             }
           in
           let expr = Closure (missing, (free_pc, missing), None) in
           update_def x expr;
-          Let (x, expr) :: acc, free_pc + 1, (free_pc, block) :: extra
-      | Some _ -> assert false)
-  | _ -> i :: acc, free_pc, extra
+          Let (x, expr) :: acc, free_pc + 1, (free_pc, block) :: extra)
+        else (
+          assert (arity < nargs);
+          (* over application *)
+          incr opt_count;
+          let v = Code.Var.fresh () in
+          let args, rest = List.take arity args in
+          let exact_expr = Apply { f; args; exact = true } in
+          let body =
+            (* Reversed *)
+            add_event loc (Let (v, exact_expr) :: acc)
+          in
+          loop x v rest res loc (body, free_pc, extra))
+  in
+  fun i (((body_rev, free_pc, extra) as acc), loc) ->
+    match i with
+    | Let (x, Apply { f; args; exact = false }) -> loop x f args (shape f) loc acc
+    | _ -> i :: body_rev, free_pc, extra
 
-let specialize_instrs ~function_arity ~update_def opt_count p =
+let specialize_instrs ~shape ~update_def opt_count p =
   let blocks, free_pc =
+    let specialize_instrs = specialize_apply opt_count shape update_def in
     Addr.Map.fold
       (fun pc block (blocks, free_pc) ->
         if List.exists ~f:unknown_apply block.body
         then
-          let (body, free_pc, extra), _ =
+          let (body_rev, free_pc, extra), _ =
             List.fold_left
               block.body
               ~init:(([], free_pc, []), None)
               ~f:(fun acc i ->
                 match i with
                 | Event loc ->
-                    let (body, free_pc, extra), _ = acc in
-                    (i :: body, free_pc, extra), Some loc
-                | _ -> specialize_apply opt_count function_arity update_def acc i, None)
+                    let (body_rev, free_pc, extra), _ = acc in
+                    (i :: body_rev, free_pc, extra), Some loc
+                | _ -> specialize_instrs i acc, None)
           in
           let blocks =
             List.fold_left extra ~init:blocks ~f:(fun blocks (pc, b) ->
                 Addr.Map.add pc b blocks)
           in
-          Addr.Map.add pc { block with Code.body = List.rev body } blocks, free_pc
+          Addr.Map.add pc { block with Code.body = List.rev body_rev } blocks, free_pc
         else blocks, free_pc)
       p.blocks
       (p.blocks, p.free_pc)
   in
   { p with blocks; free_pc }
 
-let f ~function_arity ~update_def p =
+let f ~shape ~update_def p =
   Code.invariant p;
   let previous_p = p in
   let t = Timer.make () in
   let opt_count = ref 0 in
   let p =
-    if Config.Flag.optcall ()
-    then specialize_instrs ~function_arity ~update_def opt_count p
-    else p
+    if Config.Flag.optcall () then specialize_instrs ~shape ~update_def opt_count p else p
   in
   if times () then Format.eprintf "  optcall: %a@." Timer.print t;
   if stats () then Format.eprintf "Stats - optcall: %d@." !opt_count;

compiler/lib/specialize.mli

@@ -18,11 +18,8 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-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)
+     shape:(Code.Var.t -> Shape.t)
   -> update_def:(Code.Var.t -> Code.expr -> unit)
   -> Code.program
   -> Code.program