Do not store the code pointer in the closure if we know we don't need it

Author: vouillon

compiler/lib-wasm/gc_target.ml

@@ -240,18 +240,7 @@ module Type = struct
     then
       register_type
         (if cps then "cps_closure_0" else "closure_0")
-        (fun () ->
-          let* fun_ty' = function_type ~cps arity in
-          return
-            { supertype = None
-            ; final = false
-            ; typ =
-                W.Struct
-                  [ { mut = false
-                    ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
-                    }
-                  ]
-            })
+        (fun () -> return { supertype = None; final = false; typ = W.Struct [] })
     else
       register_type
         (if cps
@@ -287,30 +276,29 @@ module Type = struct
        then Printf.sprintf "cps_env_%d_%d" arity env_type_id
        else Printf.sprintf "env_%d_%d" arity env_type_id)
       (fun () ->
-        let* cl_typ = closure_type ~usage:`Alloc ~cps arity in
-        let* common = closure_common_fields ~cps in
-        let* fun_ty' = function_type ~cps arity in
-        return
-          { supertype = Some cl_typ
-          ; final = true
-          ; typ =
-              W.Struct
-                ((if arity = 1
-                  then common
-                  else if arity = 0
-                  then
-                    [ { mut = false
-                      ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
-                      }
-                    ]
-                  else
-                    common
-                    @ [ { mut = false
-                        ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
-                        }
-                      ])
-                @ make_env_type env_type)
-          })
+        if arity = 0
+        then
+          return
+            { supertype = None; final = true; typ = W.Struct (make_env_type env_type) }
+        else
+          let* common = closure_common_fields ~cps in
+          let* cl_typ = closure_type ~usage:`Alloc ~cps arity in
+          let* fun_ty' = function_type ~cps arity in
+          return
+            { supertype = Some cl_typ
+            ; final = true
+            ; typ =
+                W.Struct
+                  ((if arity = 1
+                    then common
+                    else
+                      common
+                      @ [ { mut = false
+                          ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
+                          }
+                        ])
+                  @ make_env_type env_type)
+            })
 
   let rec_env_type ~function_count ~env_type_id ~env_type =
     register_type (Printf.sprintf "rec_env_%d_%d" function_count env_type_id) (fun () ->
@@ -334,28 +322,41 @@ module Type = struct
        then Printf.sprintf "cps_closure_rec_%d_%d_%d" arity function_count env_type_id
        else Printf.sprintf "closure_rec_%d_%d_%d" arity function_count env_type_id)
       (fun () ->
-        let* cl_typ = closure_type ~usage:`Alloc ~cps arity in
-        let* common = closure_common_fields ~cps in
-        let* fun_ty' = function_type ~cps arity in
         let* env_ty = rec_env_type ~function_count ~env_type_id ~env_type in
-        return
-          { supertype = Some cl_typ
-          ; final = true
-          ; typ =
-              W.Struct
-                ((if arity = 1
-                  then common
-                  else
-                    common
-                    @ [ { mut = false
-                        ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
-                        }
-                      ])
-                @ [ { W.mut = false
+        if arity = 0
+        then
+          return
+            { supertype = None
+            ; final = true
+            ; typ =
+                W.Struct
+                  [ { W.mut = false
                     ; typ = W.Value (Ref { nullable = false; typ = Type env_ty })
                     }
-                  ])
-          })
+                  ]
+            }
+        else
+          let* cl_typ = closure_type ~usage:`Alloc ~cps arity in
+          let* common = closure_common_fields ~cps in
+          let* fun_ty' = function_type ~cps arity in
+          return
+            { supertype = Some cl_typ
+            ; final = true
+            ; typ =
+                W.Struct
+                  ((if arity = 1
+                    then common
+                    else
+                      common
+                      @ [ { mut = false
+                          ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
+                          }
+                        ])
+                  @ [ { W.mut = false
+                      ; typ = W.Value (Ref { nullable = false; typ = Type env_ty })
+                      }
+                    ])
+            })
 
   let rec curry_type ~cps arity m =
     register_type
@@ -800,7 +801,8 @@ module Memory = struct
 
   let env_start arity =
     match arity with
-    | 0 | 1 -> 1
+    | 0 -> 0
+    | 1 -> 1
     | _ -> 2
 
   let load_function_pointer ~cps ~arity ?(skip_cast = false) closure =
@@ -1053,7 +1055,7 @@ module Closure = struct
     | [ (g, _) ] -> Code.Var.equal f g
     | _ :: r -> is_last_fun r f
 
-  let translate ~context ~closures ~cps f =
+  let translate ~context ~closures ~cps ~need_pointer f =
     let info = Code.Var.Map.find f closures in
     let free_variables = get_free_variables ~context info in
     assert (
@@ -1062,7 +1064,7 @@ module Closure = struct
            ~f:(fun x -> Code.Var.Set.mem x context.globalized_variables)
            free_variables));
     let _, arity = List.find ~f:(fun (f', _) -> Code.Var.equal f f') info.functions in
-    let arity = if cps then arity - 1 else arity in
+    let arity = if need_pointer then if cps then arity - 1 else arity else 0 in
     let* curry_fun = if arity > 1 then need_curry_fun ~cps ~arity else return f in
     if List.is_empty free_variables
     then
@@ -1075,7 +1077,8 @@ module Closure = struct
           (W.StructNew
              ( typ
              , match arity with
-               | 0 | 1 -> [ W.RefFunc f ]
+               | 0 -> []
+               | 1 -> [ W.RefFunc f ]
                | _ -> [ RefFunc curry_fun; RefFunc f ] ))
       in
       return (W.GlobalGet name)
@@ -1098,7 +1101,8 @@ module Closure = struct
             (W.StructNew
                ( typ
                , (match arity with
-                 | 0 | 1 -> [ W.RefFunc f ]
+                 | 0 -> []
+                 | 1 -> [ W.RefFunc f ]
                  | _ -> [ RefFunc curry_fun; RefFunc f ])
                  @ l ))
       | (g, _) :: _ as functions ->
@@ -1132,9 +1136,10 @@ module Closure = struct
             return
               (W.StructNew
                  ( typ
-                 , (if arity = 1
-                    then [ W.RefFunc f ]
-                    else [ RefFunc curry_fun; RefFunc f ])
+                 , (match arity with
+                   | 0 -> []
+                   | 1 -> [ W.RefFunc f ]
+                   | _ -> [ RefFunc curry_fun; RefFunc f ])
                    @ [ env ] ))
           in
           if is_last_fun functions f
@@ -1155,7 +1160,7 @@ module Closure = struct
               (load f)
           else res
 
-  let bind_environment ~context ~closures ~cps f =
+  let bind_environment ~context ~closures ~cps ~need_pointer f =
     let info = Code.Var.Map.find f closures in
     let free_variables = get_free_variables ~context info in
     let free_variable_count = List.length free_variables in
@@ -1167,7 +1172,7 @@ module Closure = struct
     else
       let env_type_id = Option.value ~default:(-1) info.id in
       let _, arity = List.find ~f:(fun (f', _) -> Code.Var.equal f f') info.functions in
-      let arity = if cps then arity - 1 else arity in
+      let arity = if need_pointer then if cps then arity - 1 else arity else 0 in
       let offset = Memory.env_start arity in
       match info.Closure_conversion.functions with
       | [ _ ] ->

compiler/lib-wasm/generate.ml

@@ -37,6 +37,7 @@ module Generate (Target : Target_sig.S) = struct
     ; in_cps : Effects.in_cps
     ; deadcode_sentinal : Var.t
     ; global_flow_info : Global_flow.info
+    ; fun_info : Call_graph_analysis.t
     ; types : Typing.typ Var.Tbl.t
     ; blocks : block Addr.Map.t
     ; closures : Closure_conversion.closure Var.Map.t
@@ -791,35 +792,40 @@ module Generate (Target : Target_sig.S) = struct
           let rec loop acc l =
             match l with
             | [] -> (
-                let arity = List.length args in
-                let funct = Var.fresh () in
-                let* closure = tee funct (load f) in
-                let* ty, funct =
-                  Memory.load_function_pointer
-                    ~cps:(Var.Set.mem x ctx.in_cps)
-                    ~arity
-                    (load funct)
-                in
                 let* b = is_closure f in
                 if b
-                then return (W.Call (f, List.rev (closure :: acc)))
+                then
+                  let* closure = load f in
+                  return (W.Call (f, List.rev (closure :: acc)))
                 else
-                  match funct 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
+                  match
+                    if exact
+                    then Global_flow.get_unique_closure ctx.global_flow_info f
+                    else None
+                  with
+                  | Some g ->
+                      let* closure = load f in
+                      let* cl =
+                        (* Functions with constant closures ignore their
+                           environment. *)
+                        match closure with
+                        | GlobalGet _ -> Value.unit
+                        | _ -> return closure
+                      in
                       return (W.Call (g, List.rev (cl :: acc)))
-                  | _ -> (
-                      match
-                        if exact
-                        then Global_flow.get_unique_closure ctx.global_flow_info f
-                        else None
-                      with
-                      | Some g -> return (W.Call (g, List.rev (closure :: acc)))
-                      | None -> return (W.Call_ref (ty, funct, List.rev (closure :: acc)))
-                      ))
+                  | None -> (
+                      let arity = List.length args in
+                      let funct = Var.fresh () in
+                      let* closure = tee funct (load f) in
+                      let* ty, funct =
+                        Memory.load_function_pointer
+                          ~cps:(Var.Set.mem x ctx.in_cps)
+                          ~arity
+                          (load funct)
+                      in
+                      match funct with
+                      | W.RefFunc g -> return (W.Call (g, List.rev (closure :: acc)))
+                      | _ -> return (W.Call_ref (ty, funct, List.rev (closure :: acc)))))
             | x :: r ->
                 let* x = load_and_box ctx x in
                 loop (x :: acc) r
@@ -848,6 +854,10 @@ module Generate (Target : Target_sig.S) = struct
           ~context:ctx.global_context
           ~closures:ctx.closures
           ~cps:(Var.Set.mem x ctx.in_cps)
+          ~need_pointer:
+            (not
+               (Config.Flag.optcall ()
+               && Var.Hashtbl.mem ctx.fun_info.unambiguous_non_escaping x))
           x
     | Constant c -> Constant.translate c
     | Special (Alias_prim _) -> assert false
@@ -1272,6 +1282,10 @@ module Generate (Target : Target_sig.S) = struct
             ~context:ctx.global_context
             ~closures:ctx.closures
             ~cps:(Var.Set.mem f ctx.in_cps)
+            ~need_pointer:
+              (not
+                 (Config.Flag.optcall ()
+                 && Var.Hashtbl.mem ctx.fun_info.unambiguous_non_escaping f))
             f
       | None -> return ()
     in
@@ -1401,6 +1415,7 @@ module Generate (Target : Target_sig.S) = struct
 *)
       ~deadcode_sentinal
       ~global_flow_info
+      ~fun_info
       ~types =
     global_context.unit_name <- unit_name;
     let p, closures = Closure_conversion.f p in
@@ -1412,6 +1427,7 @@ module Generate (Target : Target_sig.S) = struct
       ; in_cps
       ; deadcode_sentinal
       ; global_flow_info
+      ; fun_info
       ; types
       ; blocks = p.blocks
       ; closures
@@ -1521,6 +1537,7 @@ 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 t = Timer.make () in
   let p = fix_switch_branches p in
@@ -1532,6 +1549,7 @@ let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal ~global_flow_d
       ~in_cps
       ~deadcode_sentinal
       ~global_flow_info:info
+      ~fun_info
       ~types
       p
   in

compiler/lib-wasm/target_sig.ml

@@ -174,13 +174,15 @@ module type S = sig
          context:Code_generation.context
       -> closures:Closure_conversion.closure Code.Var.Map.t
       -> cps:bool
+      -> need_pointer:bool
       -> Code.Var.t
       -> expression
 
     val bind_environment :
          context:Code_generation.context
       -> closures:Closure_conversion.closure Code.Var.Map.t
       -> cps:bool
+      -> need_pointer:bool
       -> Code.Var.t
       -> unit Code_generation.t