WIP

Author: vouillon

compiler/lib-wasm/generate.ml

@@ -812,7 +812,24 @@ module Generate (Target : Target_sig.S) = struct
                         | GlobalGet _ -> Value.unit
                         | _ -> return closure
                       in
-                      return (W.Call (g, List.rev (cl :: acc)))
+                      let params =
+                        match ctx.global_flow_info.info_defs.(Var.idx g) with
+                        | Expr (Closure (params, _, _)) -> params
+                        | _ -> assert false
+                      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, List.rev (cl :: List.rev args)))
                   | None -> (
                       let arity = List.length args in
                       let funct = Var.fresh () in
@@ -1341,7 +1358,23 @@ 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
+                Var.Hashtbl.mem
+                  ctx.fun_info.Call_graph_analysis.unambiguous_non_escaping
+                  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
@@ -1536,9 +1569,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 =
@@ -1548,7 +1583,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,27 @@ 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
+                   (Var.Hashtbl.mem
+                      fun_info.Call_graph_analysis.unambiguous_non_escaping
+                      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 +324,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 +337,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 +353,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 +382,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 && Var.Hashtbl.mem st.fun_info.unambiguous_non_escaping g
+                      then res
+                      else Domain.box res
                   | Expr (Closure (_, _, _)) ->
                       (* The function is partially applied or over applied *)
                       Top
@@ -403,33 +412,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