When a function is only called directy, pass its integer parameters as i32.

Author: vouillon

compiler/lib-wasm/generate.ml

@@ -788,13 +788,12 @@ module Generate (Target : Target_sig.S) = struct
     match e with
     | Apply { f; args; exact; _ } ->
         let* closure = load f in
-        let* args = expression_list (fun x -> load_and_box ctx x) args in
         if exact || List.length args = if Var.Set.mem x ctx.in_cps then 2 else 1
         then
           match
             if exact then Global_flow.get_unique_closure ctx.global_flow_info f else None
           with
-          | Some g ->
+          | Some (g, params) ->
               let* cl =
                 (* Functions with constant closures ignore their environment. *)
                 match closure with
@@ -803,6 +802,18 @@ module Generate (Target : Target_sig.S) = struct
                     if Option.is_some init then Value.unit else return closure
                 | _ -> return closure
               in
+              let* args =
+                expression_list
+                  Fun.id
+                  (List.map2
+                     ~f:(fun a p ->
+                       convert
+                         ~from:(get_var_type ctx a)
+                         ~into:(get_var_type ctx p)
+                         (load a))
+                     args
+                     params)
+              in
               return (W.Call (g, args @ [ cl ]))
           | None -> (
               let funct = Var.fresh () in
@@ -813,13 +824,15 @@ module Generate (Target : Target_sig.S) = struct
                   ~arity:(List.length args)
                   (load funct)
               in
+              let* args = expression_list (fun x -> load_and_box ctx x) args in
               match funct with
               | W.RefFunc g -> return (W.Call (g, args @ [ closure ]))
               | _ -> return (W.Call_ref (ty, funct, args @ [ closure ])))
         else
           let* apply =
             need_apply_fun ~cps:(Var.Set.mem x ctx.in_cps) ~arity:(List.length args)
           in
+          let* args = expression_list (fun x -> load_and_box ctx x) args in
           return (W.Call (apply, args @ [ closure ]))
     | Block (tag, a, _, _) ->
         Memory.allocate
@@ -1249,7 +1262,14 @@ module Generate (Target : Target_sig.S) = struct
       List.fold_left
         ~f:(fun l x ->
           let* _ = l in
-          let* _ = add_var x in
+          let* _ =
+            add_var
+              ?typ:
+                (match get_var_type ctx x with
+                | Typing.Int (Normalized | Unnormalized) -> Some I32
+                | _ -> None)
+              x
+          in
           return ())
         ~init:(return ())
         params
@@ -1318,7 +1338,20 @@ module Generate (Target : Target_sig.S) = struct
       ; signature =
           (match name_opt with
           | None -> Type.primitive_type param_count
-          | Some _ -> Type.func_type (param_count - 1))
+          | Some f ->
+              if Call_graph_analysis.direct_calls_only ctx.fun_info f
+              then
+                { W.params =
+                    List.map
+                      ~f:(fun x : W.value_type ->
+                        match get_var_type ctx x with
+                        | Int (Unnormalized | Normalized) -> I32
+                        | _ -> Type.value)
+                      params
+                    @ [ Type.value ]
+                ; result = [ Type.value ]
+                }
+              else Type.func_type (param_count - 1))
       ; param_names
       ; locals
       ; body
@@ -1513,9 +1546,11 @@ let init = G.init
 let start () = make_context ~value_type:Gc_target.Type.value
 
 let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal ~global_flow_data =
-  let state, info = global_flow_data in
-  let fun_info = Call_graph_analysis.f p info in
-  let types = Typing.f ~state ~info ~deadcode_sentinal p in
+  let global_flow_state, global_flow_info = global_flow_data in
+  let fun_info = Call_graph_analysis.f p global_flow_info in
+  let types =
+    Typing.f ~global_flow_state ~global_flow_info ~fun_info ~deadcode_sentinal p
+  in
   let t = Timer.make () in
   let p = fix_switch_branches p in
   let res =
@@ -1525,7 +1560,7 @@ let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal ~global_flow_d
       ~live_vars
       ~in_cps
       ~deadcode_sentinal
-      ~global_flow_info:info
+      ~global_flow_info
       ~fun_info
       ~types
       p

compiler/lib-wasm/typing.ml

@@ -128,22 +128,25 @@ let update_deps st { blocks; _ } =
           | _ -> ()))
     blocks
 
-let mark_function_parameters { blocks; _ } =
-  let function_parameters = Var.Tbl.make () false in
-  let set x = Var.Tbl.set function_parameters x true in
+let mark_function_parameters ~fun_info { blocks; _ } =
+  let boxed_function_parameters = Var.Tbl.make () false in
+  let set x = Var.Tbl.set boxed_function_parameters x true in
   Addr.Map.iter
     (fun _ block ->
       List.iter block.body ~f:(fun i ->
           match i with
-          | Let (_, Closure (params, _, _)) -> List.iter ~f:set params
+          | Let (x, Closure (params, _, _))
+            when not (Call_graph_analysis.direct_calls_only fun_info x) ->
+              List.iter ~f:set params
           | _ -> ()))
     blocks;
-  function_parameters
+  boxed_function_parameters
 
 type st =
-  { state : state
-  ; info : info
-  ; function_parameters : bool Var.Tbl.t
+  { global_flow_state : state
+  ; global_flow_info : info
+  ; boxed_function_parameters : bool Var.Tbl.t
+  ; fun_info : Call_graph_analysis.t
   }
 
 let rec constant_type (c : constant) =
@@ -319,11 +322,11 @@ let prim_type ~approx prim args =
   | _ -> Top
 
 let propagate st approx x : Domain.t =
-  match st.state.defs.(Var.idx x) with
+  match st.global_flow_state.defs.(Var.idx x) with
   | Phi { known; others; unit } ->
       let res = Domain.join_set ~others (fun y -> Var.Tbl.get approx y) known in
       let res = if unit then Domain.join (Int Unnormalized) res else res in
-      if Var.Tbl.get st.function_parameters x then Domain.box res else res
+      if Var.Tbl.get st.boxed_function_parameters x then Domain.box res else res
   | Expr e -> (
       match e with
       | Constant c -> constant_type c
@@ -332,7 +335,7 @@ let propagate st approx x : Domain.t =
           Tuple
             (Array.mapi
                ~f:(fun i y ->
-                 match st.state.mutable_fields.(Var.idx x) with
+                 match st.global_flow_state.mutable_fields.(Var.idx x) with
                  | All_fields -> Top
                  | Some_fields s when IntSet.mem i s -> Top
                  | Some_fields _ | No_field ->
@@ -348,15 +351,15 @@ let propagate st approx x : Domain.t =
           ( Extern ("caml_check_bound" | "caml_check_bound_float" | "caml_check_bound_gen")
           , [ Pv y; _ ] ) -> Var.Tbl.get approx y
       | Prim ((Array_get | Extern "caml_array_unsafe_get"), [ Pv y; _ ]) -> (
-          match Var.Tbl.get st.info.info_approximation y with
+          match Var.Tbl.get st.global_flow_info.info_approximation y with
           | Values { known; others } ->
               Domain.join_set
                 ~others
                 (fun z ->
-                  match st.state.defs.(Var.idx z) with
+                  match st.global_flow_state.defs.(Var.idx z) with
                   | Expr (Block (_, lst, _, _)) ->
                       let m =
-                        match st.state.mutable_fields.(Var.idx z) with
+                        match st.global_flow_state.mutable_fields.(Var.idx z) with
                         | No_field -> false
                         | Some_fields _ | All_fields -> true
                       in
@@ -377,18 +380,22 @@ let propagate st approx x : Domain.t =
       | Prim (Extern prim, args) -> prim_type ~approx prim args
       | Special _ -> Top
       | Apply { f; args; _ } -> (
-          match Var.Tbl.get st.info.info_approximation f with
+          match Var.Tbl.get st.global_flow_info.info_approximation f with
           | Values { known; others } ->
               Domain.join_set
                 ~others
                 (fun g ->
-                  match st.state.defs.(Var.idx g) with
+                  match st.global_flow_state.defs.(Var.idx g) with
                   | Expr (Closure (params, _, _))
                     when List.length args = List.length params ->
-                      Domain.box
-                        (Domain.join_set
-                           (fun y -> Var.Tbl.get approx y)
-                           (Var.Map.find g st.state.return_values))
+                      let res =
+                        Domain.join_set
+                          (fun y -> Var.Tbl.get approx y)
+                          (Var.Map.find g st.global_flow_state.return_values)
+                      in
+                      if false && Call_graph_analysis.direct_calls_only st.fun_info g
+                      then res
+                      else Domain.box res
                   | Expr (Closure (_, _, _)) ->
                       (* The function is partially applied or over applied *)
                       Top
@@ -403,33 +410,36 @@ module Solver = G.Solver (Domain)
 let solver st =
   let associated_list h x = try Var.Hashtbl.find h x with Not_found -> [] in
   let g =
-    { G.domain = st.state.vars
+    { G.domain = st.global_flow_state.vars
     ; G.iter_children =
         (fun f x ->
-          List.iter ~f (Var.Tbl.get st.state.deps x);
+          List.iter ~f (Var.Tbl.get st.global_flow_state.deps x);
           List.iter
-            ~f:(fun g -> List.iter ~f (associated_list st.state.function_call_sites g))
-            (associated_list st.state.functions_from_returned_value x))
+            ~f:(fun g ->
+              List.iter ~f (associated_list st.global_flow_state.function_call_sites g))
+            (associated_list st.global_flow_state.functions_from_returned_value x))
     }
   in
   Solver.f () g (propagate st)
 
-let f ~state ~info ~deadcode_sentinal p =
-  update_deps state p;
-  let function_parameters = mark_function_parameters p in
-  let typ = solver { state; info; function_parameters } in
+let f ~global_flow_state ~global_flow_info ~fun_info ~deadcode_sentinal p =
+  update_deps global_flow_state p;
+  let boxed_function_parameters = mark_function_parameters ~fun_info p in
+  let typ =
+    solver { global_flow_state; global_flow_info; fun_info; boxed_function_parameters }
+  in
   Var.Tbl.set typ deadcode_sentinal (Int Normalized);
   if debug ()
   then (
     Var.ISet.iter
       (fun x ->
-        match state.defs.(Var.idx x) with
+        match global_flow_state.defs.(Var.idx x) with
         | Expr _ -> ()
         | Phi _ ->
             let t = Var.Tbl.get typ x in
             if not (Domain.equal t Top)
             then Format.eprintf "%a: %a@." Var.print x Domain.print t)
-      state.vars;
+      global_flow_state.vars;
     Print.program
       Format.err_formatter
       (fun _ i ->

compiler/lib-wasm/typing.mli

@@ -21,8 +21,9 @@ type typ =
 val constant_type : Code.constant -> typ
 
 val f :
-     state:Global_flow.state
-  -> info:Global_flow.info
+     global_flow_state:Global_flow.state
+  -> global_flow_info:Global_flow.info
+  -> fun_info:Call_graph_analysis.t
   -> deadcode_sentinal:Code.Var.t
   -> Code.program
   -> typ Code.Var.Tbl.t

compiler/lib/global_flow.ml

@@ -791,9 +791,9 @@ let get_unique_closure info f =
           Var.Set.fold
             (fun g acc ->
               match info.info_defs.(Var.idx g) with
-              | Expr (Closure _) -> (
+              | Expr (Closure (params, _, _)) -> (
                   match acc with
-                  | None -> Some (Some g)
+                  | None -> Some (Some (g, params))
                   | Some (Some _) -> Some None
                   | Some None -> acc)
               | Expr (Block _) -> acc

compiler/lib/global_flow.mli

@@ -84,6 +84,6 @@ val update_def : info -> Code.Var.t -> Code.expr -> unit
 
 val exact_call : info -> Var.t -> int -> bool
 
-val get_unique_closure : info -> Var.t -> Var.t option
+val get_unique_closure : info -> Var.t -> (Var.t * Var.t list) option
 
 val function_arity : info -> Var.t -> int option