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

@@ -510,6 +510,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 = 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 rec store ?(always = false) ?typ x e =

compiler/lib-wasm/code_generation.mli

@@ -203,3 +203,5 @@ val function_body :
   -> param_names:Code.Var.t list
   -> body:unit t
   -> (Wasm_ast.var * Wasm_ast.value_type) list * Wasm_ast.instruction list
+
+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 =
@@ -119,7 +124,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
@@ -140,7 +145,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 =
@@ -193,7 +198,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))
@@ -302,7 +307,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

@@ -27,6 +27,8 @@ let include_closure_arity = false
 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
@@ -205,7 +207,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
@@ -433,6 +436,8 @@ module Value = struct
     let* t = Type.block_type in
     return (W.ArrayNewFixed (t, []))
 
+  let dummy_closure = empty_struct
+
   let as_block e =
     let* t = Type.block_type in
     let* e = e in
@@ -816,6 +821,11 @@ module Memory = struct
     then 1
     else (if include_closure_arity then 1 else 0) + if arity = 1 then 1 else 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 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
@@ -1192,7 +1202,7 @@ module Closure = struct
     if free_variable_count = 0
     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 arity = List.assoc f info.functions in
@@ -1201,7 +1211,7 @@ module Closure = struct
       match info.Closure_conversion.functions with
       | [ _ ] ->
           let* typ = Type.env_type ~cps ~arity free_variable_count in
-          let* _ = add_var f in
+          let* _ = add_var ~typ:Type.closure f in
           let env = Code.Var.fresh_n "env" in
           let* () =
             store
@@ -1222,7 +1232,7 @@ module Closure = struct
           let* typ =
             Type.rec_closure_type ~cps ~arity ~function_count ~free_variable_count
           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 ~free_variable_count in
           let* () =

compiler/lib-wasm/generate.ml

@@ -192,7 +192,12 @@ module Generate (Target : Target_sig.S) = struct
           | [] -> (
               let arity = List.length args in
               let funct = Var.fresh () in
-              let* closure = tee funct (load f) in
+              let* closure =
+                Memory.cast_closure
+                  ~cps:(Var.Set.mem x ctx.in_cps)
+                  ~arity
+                  (tee funct (load f))
+              in
               let* ty, funct =
                 Memory.load_function_pointer
                   ~cps:(Var.Set.mem x ctx.in_cps)
@@ -208,7 +213,7 @@ module Generate (Target : Target_sig.S) = struct
                     (* 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 then Value.dummy_closure else return closure in
                     return (W.Call (g, List.rev (cl :: acc)))
                 | _ -> return (W.Call_ref (ty, funct, List.rev (closure :: acc))))
           | x :: r ->

compiler/lib-wasm/target_sig.ml

@@ -26,6 +26,8 @@ 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 load_function_pointer :
          cps:bool
       -> arity:int
@@ -99,6 +101,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
@@ -159,6 +163,8 @@ module type S = sig
 
     val dummy_block : expression
 
+    val dummy_closure : expression
+
     val as_block : expression -> expression
   end