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

@@ -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
@@ -430,6 +433,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
@@ -811,6 +816,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
@@ -1189,7 +1199,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 env_type_id = Option.value ~default:(-1) info.id in
@@ -1199,7 +1209,7 @@ module Closure = struct
       match info.Closure_conversion.functions with
       | [ _ ] ->
           let* typ = Type.env_type ~cps ~arity ~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
@@ -1220,7 +1230,7 @@ module Closure = struct
           let* typ =
             Type.rec_closure_type ~cps ~arity ~function_count ~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

@@ -793,7 +793,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)
@@ -809,7 +814,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)))
                   | _ -> (
                       match

compiler/lib-wasm/target_sig.ml

@@ -27,6 +27,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
@@ -100,6 +102,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 +166,8 @@ module type S = sig
 
     val dummy_block : expression
 
+    val dummy_closure : expression
+
     val as_block : expression -> expression
   end