Optimize calling a known function

Author: vouillon

compiler/lib-wasm/generate.ml

@@ -178,10 +178,16 @@ module Generate (Target : Target_sig.S) = struct
 
   let zero_divide_pc = -2
 
+  let exact_call kind =
+    match kind with
+    | Generic -> false
+    | Exact | Known _ -> true
+
   let rec translate_expr ctx context x e =
     match e with
-    | Apply { f; args; exact }
-      when exact || List.length args = if Var.Set.mem x ctx.in_cps then 2 else 1 ->
+    | Apply { f; args; kind }
+      when exact_call kind || List.length args = if Var.Set.mem x ctx.in_cps then 2 else 1
+      ->
         let rec loop acc l =
           match l with
           | [] -> (
@@ -198,13 +204,14 @@ module Generate (Target : Target_sig.S) = struct
               if b
               then return (W.Call (f, List.rev (closure :: acc)))
               else
-                match funct with
-                | W.RefFunc g ->
+                match funct, kind with
+                | W.RefFunc g, _ ->
                     (* Functions with constant closures ignore their
                        environment. In case of partial application, we
                        still need the closure. *)
-                    let* cl = if exact then Value.unit else return closure in
+                    let* cl = if exact_call kind then Value.unit else return closure in
                     return (W.Call (g, List.rev (cl :: acc)))
+                | _, Known g -> return (W.Call (g, List.rev (closure :: acc)))
                 | _ -> return (W.Call_ref (ty, funct, List.rev (closure :: acc))))
           | x :: r ->
               let* x = load x in

compiler/lib/code.ml

@@ -408,11 +408,16 @@ type field_type =
   | Non_float
   | Float
 
+type apply_kind =
+  | Generic
+  | Exact
+  | Known of Var.t
+
 type expr =
   | Apply of
       { f : Var.t
       ; args : Var.t list
-      ; exact : bool
+      ; kind : apply_kind
       }
   | Block of int * Var.t array * array_or_not * mutability
   | Field of Var.t * int * field_type
@@ -552,10 +557,12 @@ module Print = struct
 
   let expr f e =
     match e with
-    | Apply { f = g; args; exact } ->
-        if exact
-        then Format.fprintf f "%a!(%a)" Var.print g var_list args
-        else Format.fprintf f "%a(%a)" Var.print g var_list args
+    | Apply { f = g; args; kind } -> (
+        match kind with
+        | Generic -> Format.fprintf f "%a(%a)" Var.print g var_list args
+        | Exact -> Format.fprintf f "%a!(%a)" Var.print g var_list args
+        | Known h -> Format.fprintf f "%a{=%a}(%a)" Var.print g Var.print h var_list args
+        )
     | Block (t, a, _, mut) ->
         Format.fprintf
           f

compiler/lib/code.mli

@@ -206,11 +206,16 @@ type field_type =
   | Non_float
   | Float
 
+type apply_kind =
+  | Generic
+  | Exact (* # of arguments = # of parameters *)
+  | Known of Var.t (* Exact and we know which function is called *)
+
 type expr =
   | Apply of
       { f : Var.t
       ; args : Var.t list
-      ; exact : bool (* if true, then # of arguments = # of parameters *)
+      ; kind : apply_kind
       }
   | Block of int * Var.t array * array_or_not * mutability
   | Field of Var.t * int * field_type

compiler/lib/duplicate.ml

@@ -24,8 +24,10 @@ let subst_cont m s (pc, arg) = Addr.Map.find pc m, List.map arg ~f:(fun x -> s x
 let expr s e =
   match e with
   | Constant _ -> e
-  | Apply { f; args; exact } ->
-      Apply { f = s f; args = List.map args ~f:(fun x -> s x); exact }
+  | Apply { f; args; kind = Known g } ->
+      Apply { f = s f; args = List.map args ~f:(fun x -> s x); kind = Known (s g) }
+  | Apply { f; args; kind = (Generic | Exact) as kind } ->
+      Apply { f = s f; args = List.map args ~f:(fun x -> s x); kind }
   | Block (n, a, k, mut) -> Block (n, Array.map a ~f:(fun x -> s x), k, mut)
   | Field (x, n, field_type) -> Field (s x, n, field_type)
   | Closure _ -> failwith "Inlining/Duplicating closure is currenly not supported"

compiler/lib/effects.ml

@@ -304,12 +304,15 @@ let allocate_closure ~st ~params ~body ~branch =
   let name = Var.fresh () in
   [ Let (name, Closure (params, (pc, []))) ], name
 
-let tail_call ~st ?(instrs = []) ~exact ~in_cps ~check ~f args =
-  assert (exact || check);
+let tail_call ~st ?(instrs = []) ~kind ~in_cps ~check ~f args =
+  assert (
+    match kind with
+    | Generic -> check
+    | Exact | Known _ -> true);
   let ret = Var.fresh () in
   if check then st.trampolined_calls := Var.Set.add ret !(st.trampolined_calls);
   if in_cps then st.in_cps := Var.Set.add ret !(st.in_cps);
-  instrs @ [ Let (ret, Apply { f; args; exact }) ], Return ret
+  instrs @ [ Let (ret, Apply { f; args; kind }) ], Return ret
 
 let cps_branch ~st ~src (pc, args) =
   match Addr.Set.mem pc st.blocks_to_transform with
@@ -327,14 +330,8 @@ let cps_branch ~st ~src (pc, args) =
       (* We check the stack depth only for backward edges (so, at
          least once per loop iteration) *)
       let check = Hashtbl.find st.block_order src >= Hashtbl.find st.block_order pc in
-      tail_call
-        ~st
-        ~instrs
-        ~exact:true
-        ~in_cps:false
-        ~check
-        ~f:(closure_of_pc ~st pc)
-        args
+      let f = closure_of_pc ~st pc in
+      tail_call ~st ~instrs ~kind:(Known f) ~in_cps:false ~check ~f args
 
 let cps_jump_cont ~st ~src ((pc, _) as cont) =
   match Addr.Set.mem pc st.blocks_to_transform with
@@ -396,7 +393,7 @@ let cps_last ~st ~alloc_jump_closures pc (last : last) ~k : instr list * last =
       (* Is the number of successive 'returns' is unbounded is CPS, it
          means that we have an unbounded of calls in direct style
          (even with tail call optimization) *)
-      tail_call ~st ~exact:true ~in_cps:false ~check:false ~f:k [ x ]
+      tail_call ~st ~kind:Exact ~in_cps:false ~check:false ~f:k [ x ]
   | Raise (x, rmode) -> (
       assert (List.is_empty alloc_jump_closures);
       match Hashtbl.find_opt st.matching_exn_handler pc with
@@ -431,7 +428,7 @@ let cps_last ~st ~alloc_jump_closures pc (last : last) ~k : instr list * last =
           tail_call
             ~st
             ~instrs:(Let (exn_handler, Prim (Extern "caml_pop_trap", [])) :: instrs)
-            ~exact:true
+            ~kind:Exact
             ~in_cps:false
             ~check:false
             ~f:exn_handler
@@ -481,6 +478,14 @@ let cps_last ~st ~alloc_jump_closures pc (last : last) ~k : instr list * last =
             @ (Let (exn_handler, Prim (Extern "caml_pop_trap", [])) :: body)
           , branch ))
 
+let refine_kind k k' =
+  match k, k' with
+  | Known _, _ -> k
+  | _, Known _ -> k'
+  | Exact, _ -> k
+  | _, Exact -> k'
+  | Generic, Generic -> k
+
 let cps_instr ~st (instr : instr) : instr =
   match instr with
   | Let (x, Closure (params, (pc, _))) when Var.Set.mem x st.cps_needed ->
@@ -498,11 +503,15 @@ let cps_instr ~st (instr : instr) : instr =
                 (Extern "caml_alloc_dummy_function", [ size; Pc (Int (Targetint.succ a)) ])
             )
       | _ -> assert false)
-  | Let (x, Apply { f; args; _ }) when not (Var.Set.mem x st.cps_needed) ->
+  | Let (x, Apply { f; args; kind }) when not (Var.Set.mem x st.cps_needed) ->
       (* At the moment, we turn into CPS any function not called with
          the right number of parameter *)
-      assert (Global_flow.exact_call st.flow_info f (List.length args));
-      Let (x, Apply { f; args; exact = true })
+      let kind' = Global_flow.apply_kind st.flow_info f (List.length args) in
+      assert (
+        match kind' with
+        | Generic -> false
+        | Exact | Known _ -> true);
+      Let (x, Apply { f; args; kind = refine_kind kind kind' })
   | Let (_, (Apply _ | Prim (Extern ("%resume" | "%perform" | "%reperform"), _))) ->
       assert false
   | _ -> instr
@@ -552,21 +561,21 @@ let cps_block ~st ~k pc block =
           [ Let (x, e) ], Return x)
     in
     match e with
-    | Apply { f; args; exact } when Var.Set.mem x st.cps_needed ->
+    | Apply { f; args; kind } when Var.Set.mem x st.cps_needed ->
         Some
           (fun ~k ->
-            let exact =
-              exact || Global_flow.exact_call st.flow_info f (List.length args)
+            let kind =
+              refine_kind kind (Global_flow.apply_kind st.flow_info f (List.length args))
             in
-            tail_call ~st ~exact ~in_cps:true ~check:true ~f (args @ [ k ]))
+            tail_call ~st ~kind ~in_cps:true ~check:true ~f (args @ [ k ]))
     | Prim (Extern "%resume", [ Pv stack; Pv f; Pv arg ]) ->
         Some
           (fun ~k ->
             let k' = Var.fresh_n "cont" in
             tail_call
               ~st
               ~instrs:[ Let (k', Prim (Extern "caml_resume_stack", [ Pv stack; Pv k ])) ]
-              ~exact:(Global_flow.exact_call st.flow_info f 1)
+              ~kind:(Global_flow.apply_kind st.flow_info f 1)
               ~in_cps:true
               ~check:true
               ~f

compiler/lib/generate.ml

@@ -149,8 +149,13 @@ module Share = struct
           List.fold_left block.body ~init:share ~f:(fun share i ->
               match i with
               | Let (_, Constant c) -> get_constant c share
-              | Let (x, Apply { args; exact; _ }) ->
+              | Let (x, Apply { args; kind; _ }) ->
                   let trampolined = Var.Set.mem x trampolined_calls in
+                  let exact =
+                    match kind with
+                    | Generic -> false
+                    | Exact | Known _ -> true
+                  in
                   if (not exact) || trampolined
                   then add_apply { arity = List.length args; exact; trampolined } share
                   else share
@@ -1181,7 +1186,12 @@ let remove_unused_tail_args ctx exact trampolined args =
 let rec translate_expr ctx loc x e level : (_ * J.statement_list) Expr_builder.t =
   let open Expr_builder in
   match e with
-  | Apply { f; args; exact } ->
+  | Apply { f; args; kind } ->
+      let exact =
+        match kind with
+        | Generic -> false
+        | Exact | Known _ -> true
+      in
       let trampolined = Var.Set.mem x ctx.Ctx.trampolined_calls in
       let args = remove_unused_tail_args ctx exact trampolined args in
       let* () = info ~need_loc:true mutator_p in

compiler/lib/generate_closure.ml

@@ -46,8 +46,8 @@ let rec collect_apply pc blocks visited tc =
       match block.branch with
       | Return x -> (
           match List.last block.body with
-          | Some (Let (y, Apply { f; exact = true; _ })) when Code.Var.compare x y = 0 ->
-              Some (add_multi f pc tc)
+          | Some (Let (y, Apply { f; kind = Exact | Known _; _ }))
+            when Code.Var.compare x y = 0 -> Some (add_multi f pc tc)
           | None -> None
           | Some _ -> None)
       | _ -> None
@@ -100,7 +100,7 @@ module Trampoline = struct
     match counter with
     | None ->
         { params = []
-        ; body = [ Let (return, Apply { f; args; exact = true }) ]
+        ; body = [ Let (return, Apply { f; args; kind = Known f }) ]
         ; branch = Return return
         }
     | Some counter ->
@@ -110,7 +110,7 @@ module Trampoline = struct
             [ Let
                 ( counter_plus_1
                 , Prim (Extern "%int_add", [ Pv counter; Pc (Int Targetint.one) ]) )
-            ; Let (return, Apply { f; args = counter_plus_1 :: args; exact = true })
+            ; Let (return, Apply { f; args = counter_plus_1 :: args; kind = Known f })
             ]
         ; branch = Return return
         }
@@ -139,14 +139,14 @@ module Trampoline = struct
           (match counter with
           | None ->
               [ Event loc
-              ; Let (result1, Apply { f; args; exact = true })
+              ; Let (result1, Apply { f; args; kind = Known f })
               ; Event Parse_info.zero
               ; Let (result2, Prim (Extern "caml_trampoline", [ Pv result1 ]))
               ]
           | Some counter ->
               [ Event loc
               ; Let (counter, Constant (Int Targetint.zero))
-              ; Let (result1, Apply { f; args = counter :: args; exact = true })
+              ; Let (result1, Apply { f; args = counter :: args; kind = Known f })
               ; Event Parse_info.zero
               ; Let (result2, Prim (Extern "caml_trampoline", [ Pv result1 ]))
               ])
@@ -222,7 +222,7 @@ module Trampoline = struct
                     let bounce_call_pc = free_pc + 1 in
                     let free_pc = free_pc + 2 in
                     match List.rev block.body with
-                    | Let (x, Apply { f; args; exact = true }) :: rem_rev ->
+                    | Let (x, Apply { f; args; kind = Exact | Known _ }) :: rem_rev ->
                         assert (Var.equal f ci.f_name);
                         let blocks =
                           Addr.Map.add

compiler/lib/global_flow.ml

@@ -682,17 +682,29 @@ let f ~fast p =
   ; info_return_vals = rets
   }
 
-let exact_call info f n =
+let apply_kind info f n =
   match Var.Tbl.get info.info_approximation f with
-  | Top | Values { others = true; _ } -> false
-  | Values { known; others = false } ->
-      Var.Set.for_all
-        (fun g ->
-          match info.info_defs.(Var.idx g) with
-          | Expr (Closure (params, _)) -> List.length params = n
-          | Expr (Block _) -> true
-          | Expr _ | Phi _ -> assert false)
-        known
+  | Top | Values { others = true; _ } -> Generic
+  | Values { known; others = false } -> (
+      match
+        Var.Set.fold
+          (fun g acc ->
+            match info.info_defs.(Var.idx g) with
+            | Expr (Closure (params, _)) ->
+                if List.length params = n
+                then
+                  match acc with
+                  | None -> Some (Known g)
+                  | Some (Known _) -> Some Exact
+                  | Some (Exact | Generic) -> acc
+                else Some Generic
+            | Expr (Block _) -> acc
+            | Expr _ | Phi _ -> assert false)
+          known
+          None
+      with
+      | None -> Exact
+      | Some kind -> kind)
 
 let function_arity info f =
   match Var.Tbl.get info.info_approximation f with
@@ -705,9 +717,10 @@ let function_arity info f =
             | 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)
+                | None -> Some (Some (n, Known g))
+                | Some (Some (n', _)) when n <> n' -> Some None
+                | Some (Some (_, Known _)) -> Some (Some (n, Exact))
+                | Some (None | Some (_, (Exact | Generic))) -> acc)
             | Expr (Block _) -> acc
             | Expr _ | Phi _ -> assert false)
           known

compiler/lib/global_flow.mli

@@ -46,6 +46,6 @@ type info =
 
 val f : fast:bool -> Code.program -> info
 
-val exact_call : info -> Var.t -> int -> bool
+val apply_kind : info -> Var.t -> int -> Code.apply_kind
 
-val function_arity : info -> Var.t -> int option
+val function_arity : info -> Var.t -> (int * Code.apply_kind) option

compiler/lib/inline.ml

@@ -180,7 +180,8 @@ let inline ~first_class_primitives live_vars closures name pc (outer, p) =
       ~init:([], (outer, block.branch, p))
       ~f:(fun i (rem, state) ->
         match i with
-        | Let (x, Apply { f; args; exact = true; _ }) when Var.Map.mem f closures -> (
+        | Let (x, Apply { f; args; kind = Exact | Known _; _ })
+          when Var.Map.mem f closures -> (
             let outer, branch, p = state in
             let { cl_params = params
                 ; cl_cont = clos_cont
@@ -268,7 +269,7 @@ let inline ~first_class_primitives live_vars closures name pc (outer, p) =
                   if recursive
                   then
                     ( Let (f, Closure (params, clos_cont))
-                      :: Let (x, Apply { f; args; exact = true })
+                      :: Let (x, Apply { f; args; kind = Known f })
                       :: rem
                     , (outer, branch, p) )
                   else