More precise environment type

Casting to the exact type should be slightly faster since we avoid the
test for i31.

Author: vouillon

compiler/lib-wasm/code_generation.ml

@@ -503,6 +503,31 @@ let tee ?typ x e =
 
 let should_make_global x st = Var.Set.mem x st.context.globalized_variables, st
 
+let get_constant x st = Var.Hashtbl.find_opt st.context.constants x, st
+
+let placeholder_value typ f =
+  let* c = get_constant typ in
+  match c with
+  | None ->
+      let x = Var.fresh () in
+      let* () = register_constant typ (W.GlobalGet x) in
+      let* () =
+        register_global
+          ~constant:true
+          x
+          { mut = false; typ = Ref { nullable = false; typ = Type typ } }
+          (f typ)
+      in
+      return (W.GlobalGet x)
+  | Some c -> return c
+
+let empty_struct =
+  let* typ =
+    register_type "empty_struct" (fun () ->
+        return { supertype = None; final = true; typ = W.Struct [] })
+  in
+  placeholder_value typ (fun typ -> W.StructNew (typ, []))
+
 let value_type st = st.context.value_type, st
 
 let get_constant x st = Var.Hashtbl.find_opt st.context.constants x, st

compiler/lib-wasm/code_generation.mli

@@ -205,3 +205,5 @@ val array_placeholder : Code.Var.t -> expression
 val default_value :
      Wasm_ast.value_type
   -> (Wasm_ast.expression * Wasm_ast.value_type * Wasm_ast.ref_type option) t
+
+val empty_struct : expression

compiler/lib-wasm/curry.ml

@@ -35,7 +35,12 @@ module Make (Target : Target_sig.S) = struct
 
   let call ?typ ~cps ~arity closure args =
     let funct = Var.fresh () in
-    let* closure = tee ?typ funct closure in
+    let closure = tee ?typ funct closure in
+    let* closure =
+      match typ with
+      | None -> Memory.cast_closure ~cps ~arity closure
+      | Some _ -> closure
+    in
     let args = args @ [ closure ] in
     let* ty, funct =
       Memory.load_function_pointer
@@ -68,7 +73,7 @@ module Make (Target : Target_sig.S) = struct
     let body =
       let* () = no_event in
       let* () = bind_parameters args in
-      let* _ = add_var f in
+      let* _ = add_var ~typ:Type.closure f in
       let* args' = expression_list load args in
       let* _f = load f in
       let rec loop m args closure closure_typ =
@@ -126,7 +131,7 @@ module Make (Target : Target_sig.S) = struct
     let body =
       let* () = no_event in
       let* _ = add_var x in
-      let* _ = add_var f in
+      let* _ = add_var ~typ:Type.closure f in
       push (Closure.curry_allocate ~cps:false ~arity m ~f:name' ~closure:f ~arg:x)
     in
     let param_names = [ x; f ] in
@@ -154,7 +159,7 @@ module Make (Target : Target_sig.S) = struct
     let body =
       let* () = no_event in
       let* () = bind_parameters args in
-      let* _ = add_var f in
+      let* _ = add_var ~typ:Type.closure f in
       let* args' = expression_list load args in
       let* _f = load f in
       let rec loop m args closure closure_typ =
@@ -214,7 +219,7 @@ module Make (Target : Target_sig.S) = struct
       let* () = no_event in
       let* _ = add_var x in
       let* _ = add_var cont in
-      let* _ = add_var f in
+      let* _ = add_var ~typ:Type.closure f in
       let* e = Closure.curry_allocate ~cps:true ~arity m ~f:name' ~closure:f ~arg:x in
       let* c = call ~cps:false ~arity:1 (load cont) [ e ] in
       instr (W.Return (Some c))
@@ -333,7 +338,7 @@ module Make (Target : Target_sig.S) = struct
     let body =
       let* () = no_event in
       let* () = bind_parameters l in
-      let* _ = add_var f in
+      let* _ = add_var ~typ:Type.closure f in
       let* typ, closure = Memory.load_real_closure ~cps ~arity (load f) in
       let* l = expression_list load l in
       let* e =

compiler/lib-wasm/gc_target.ml

@@ -25,6 +25,8 @@ type expression = Wasm_ast.expression Code_generation.t
 module Type = struct
   let value = W.Ref { nullable = false; typ = Eq }
 
+  let closure = W.Ref { nullable = false; typ = Struct }
+
   let block_type =
     register_type "block" (fun () ->
         return
@@ -203,7 +205,8 @@ module Type = struct
   let primitive_type n =
     { W.params = List.init ~len:n ~f:(fun _ -> value); result = [ value ] }
 
-  let func_type n = primitive_type (n + 1)
+  let func_type n =
+    { W.params = List.init ~len:n ~f:(fun _ -> value) @ [ closure ]; result = [ value ] }
 
   let function_type ~cps n =
     let n = if cps then n + 1 else n in
@@ -327,7 +330,7 @@ module Type = struct
                 (List.init
                    ~f:(fun i ->
                      { W.mut = i < function_count
-                     ; typ = W.Value (Ref { nullable = false; typ = Eq })
+                     ; typ = W.Value (Ref { nullable = false; typ = Struct })
                      })
                    ~len:function_count
                 @ make_env_type env_type)
@@ -441,6 +444,8 @@ module Value = struct
     let* t = Type.block_type in
     array_placeholder t
 
+  let dummy_closure = empty_struct
+
   let as_block e =
     let* t = Type.block_type in
     let* e = e in
@@ -825,6 +830,15 @@ module Memory = struct
       | 0 | 1 -> 1
       | _ -> 2
 
+  let cast_closure ~cps ~arity closure =
+    let arity = if cps then arity - 1 else arity in
+    let* ty = Type.closure_type ~usage:`Access ~cps arity in
+    wasm_cast ty closure
+
+  let cast_generic_closure closure =
+    let* e = closure in
+    return (W.RefCast ({ nullable = false; typ = Struct }, e))
+
   let load_function_pointer ~cps ~arity ?(skip_cast = false) closure =
     let arity = if cps then arity - 1 else arity in
     let* ty = Type.closure_type ~usage:`Access ~cps arity in
@@ -1096,7 +1110,7 @@ module Closure = struct
     if List.is_empty free_variables
     then
       if no_code_pointer
-      then Value.unit
+      then Value.dummy_closure
       else
         let* typ = Type.closure_type ~usage:`Alloc ~cps arity in
         let name = Code.Var.fork f in
@@ -1151,12 +1165,11 @@ module Closure = struct
               tee
                 ~typ:(W.Ref { nullable = false; typ = Type env_typ })
                 env
-                (return
+                (let* dummy_closure = Value.dummy_closure in
+                 return
                    (W.StructNew
                       ( env_typ
-                      , List.init ~len:function_count ~f:(fun _ ->
-                            W.RefI31 (W.Const (I32 0l)))
-                        @ l )))
+                      , List.init ~len:function_count ~f:(fun _ -> dummy_closure) @ l )))
             else
               let* env = get_closure_env g in
               let* () = set_closure_env f env in
@@ -1208,7 +1221,7 @@ module Closure = struct
     if List.is_empty free_variables
     then
       (* The closures are all constants and the environment is empty. *)
-      let* _ = add_var (Code.Var.fresh ()) in
+      let* _ = add_var ~typ:Type.closure (Code.Var.fresh ()) in
       return ()
     else
       let env_type_id = Option.value ~default:(-1) info.id in
@@ -1220,7 +1233,7 @@ module Closure = struct
           let* typ =
             Type.env_type ~cps ~arity ~no_code_pointer ~env_type_id ~env_type:[]
           in
-          let* _ = add_var f in
+          let* _ = add_var ~typ:Type.closure f in
           let env = Code.Var.fresh_n "env" in
           let* () =
             store
@@ -1247,7 +1260,7 @@ module Closure = struct
               ~env_type_id
               ~env_type:[]
           in
-          let* _ = add_var f in
+          let* _ = add_var ~typ:Type.closure f in
           let env = Code.Var.fresh_n "env" in
           let* env_typ = Type.rec_env_type ~function_count ~env_type_id ~env_type:[] in
           let* () =

compiler/lib-wasm/generate.ml

@@ -800,17 +800,23 @@ module Generate (Target : Target_sig.S) = struct
                 match closure with
                 | GlobalGet global ->
                     let* init = get_global global in
-                    if Option.is_some init then Value.unit else return closure
-                | _ -> return closure
+                    if Option.is_some init then Value.dummy_closure else return closure
+                | _ -> Memory.cast_generic_closure (return closure)
               in
               return (W.Call (g, args @ [ cl ]))
           | None -> (
               let funct = Var.fresh () in
-              let* closure = tee funct (return closure) in
+              let arity = List.length args in
+              let* closure =
+                Memory.cast_closure
+                  ~cps:(Var.Set.mem x ctx.in_cps)
+                  ~arity
+                  (tee funct (return closure))
+              in
               let* ty, funct =
                 Memory.load_function_pointer
                   ~cps:(Var.Set.mem x ctx.in_cps)
-                  ~arity:(List.length args)
+                  ~arity
                   (load funct)
               in
               match funct with

compiler/lib-wasm/target_sig.ml

@@ -27,6 +27,10 @@ module type S = sig
       -> [ `Expr of Wasm_ast.expression | `Var of Wasm_ast.var ] list
       -> expression
 
+    val cast_closure : cps:bool -> arity:int -> expression -> expression
+
+    val cast_generic_closure : expression -> expression
+
     val load_function_pointer :
          cps:bool
       -> arity:int
@@ -100,6 +104,8 @@ module type S = sig
   module Type : sig
     val value : Wasm_ast.value_type
 
+    val closure : Wasm_ast.value_type
+
     val func_type : int -> Wasm_ast.func_type
 
     val primitive_type : int -> Wasm_ast.func_type
@@ -162,6 +168,8 @@ module type S = sig
 
     val dummy_block : expression
 
+    val dummy_closure : expression
+
     val as_block : expression -> expression
   end