Support values with arbitrary types in function environments

Author: vouillon

compiler/lib-wasm/closure_conversion.ml

@@ -22,6 +22,7 @@ open Code
 type closure =
   { functions : (Var.t * int) list
   ; free_variables : Var.t list
+  ; mutable id : int option
   }
 
 module SCC = Strongly_connected_components.Make (Var)
@@ -144,7 +145,8 @@ let rec traverse var_depth closures program pc depth =
                   in
                   List.iter
                     ~f:(fun (f, _) ->
-                      closures := Var.Map.add f { functions; free_variables } !closures)
+                      closures :=
+                        Var.Map.add f { functions; free_variables; id = None } !closures)
                     functions;
                   fun_lst)
                 components

compiler/lib-wasm/closure_conversion.mli

@@ -19,6 +19,7 @@
 type closure =
   { functions : (Code.Var.t * int) list
   ; free_variables : Code.Var.t list
+  ; mutable id : int option
   }
 
 val f : Code.program -> Code.program * closure Code.Var.Map.t

compiler/lib-wasm/code_generation.ml

@@ -34,6 +34,7 @@ https://github.com/llvm/llvm-project/issues/58438
 type constant_global =
   { init : W.expression option
   ; constant : bool
+  ; typ : W.value_type
   }
 
 type context =
@@ -46,6 +47,7 @@ type context =
   ; types : Wasm_ast.type_field Var.Hashtbl.t
   ; mutable closure_envs : Var.t Var.Map.t
         (** GC: mapping of recursive functions to their shared environment *)
+  ; closure_types : (W.value_type option list, int) Hashtbl.t
   ; mutable apply_funs : Var.t IntMap.t
   ; mutable cps_apply_funs : Var.t IntMap.t
   ; mutable curry_funs : Var.t IntMap.t
@@ -68,6 +70,7 @@ let make_context ~value_type =
   ; type_names = String.Hashtbl.create 128
   ; types = Var.Hashtbl.create 128
   ; closure_envs = Var.Map.empty
+  ; closure_types = Poly.Hashtbl.create 128
   ; apply_funs = IntMap.empty
   ; cps_apply_funs = IntMap.empty
   ; curry_funs = IntMap.empty
@@ -198,6 +201,7 @@ let register_global name ?exported_name ?(constant = false) typ init st =
       name
       { init = (if not typ.mut then Some init else None)
       ; constant = (not typ.mut) || constant
+      ; typ = typ.typ
       }
       st.context.constant_globals;
   (), st
@@ -484,6 +488,68 @@ let load x =
   | Local (_, x, _) -> return (W.LocalGet x)
   | Expr e -> e
 
+let rec variable_type x st =
+  match Var.Map.find_opt x st.vars with
+  | Some (Local (_, _, typ)) -> typ, st
+  | Some (Expr e) ->
+      (let* e = e in
+       expression_type e)
+        st
+  | None -> None, st
+
+and expression_type (e : W.expression) st =
+  match e with
+  | Const _
+  | UnOp _
+  | BinOp _
+  | I32WrapI64 _
+  | I64ExtendI32 _
+  | F32DemoteF64 _
+  | F64PromoteF32 _
+  | BlockExpr _
+  | Call _
+  | RefFunc _
+  | Call_ref _
+  | I31Get _
+  | ArrayGet _
+  | ArrayLen _
+  | RefTest _
+  | RefEq _
+  | RefNull _
+  | Try _
+  | Br_on_null _ -> None, st
+  | LocalGet x | LocalTee (x, _) -> variable_type x st
+  | GlobalGet x ->
+      ( (try
+           let typ = (Var.Map.find x st.context.constant_globals).typ in
+           if Poly.equal typ st.context.value_type
+           then None
+           else
+             Some
+               (match typ with
+               | Ref { typ; nullable = true } -> Ref { typ; nullable = false }
+               | _ -> typ)
+         with Not_found -> None)
+      , st )
+  | Seq (_, e') -> expression_type e' st
+  | Pop typ -> Some typ, st
+  | RefI31 _ -> Some (Ref { nullable = false; typ = I31 }), st
+  | ArrayNew (ty, _, _)
+  | ArrayNewFixed (ty, _)
+  | ArrayNewData (ty, _, _, _)
+  | StructNew (ty, _) -> Some (Ref { nullable = false; typ = Type ty }), st
+  | StructGet (_, ty, i, _) -> (
+      match (Var.Hashtbl.find st.context.types ty).typ with
+      | Struct l -> (
+          match (List.nth l i).typ with
+          | Value typ ->
+              (if Poly.equal typ st.context.value_type then None else Some typ), st
+          | Packed _ -> assert false)
+      | Array _ | Func _ -> assert false)
+  | RefCast (typ, _) | Br_on_cast (_, _, typ, _) | Br_on_cast_fail (_, typ, _, _) ->
+      Some (Ref typ), st
+  | IfExpr (_, _, _, _) | ExternConvertAny _ | AnyConvertExtern _ -> None, st
+
 let tee ?typ x e =
   let* e = e in
   let* b = is_small_constant e in

compiler/lib-wasm/code_generation.mli

@@ -30,6 +30,7 @@ type context =
   ; types : Wasm_ast.type_field Code.Var.Hashtbl.t
   ; mutable closure_envs : Code.Var.t Code.Var.Map.t
         (** GC: mapping of recursive functions to their shared environment *)
+  ; closure_types : (Wasm_ast.value_type option list, int) Hashtbl.t
   ; mutable apply_funs : Code.Var.t Stdlib.IntMap.t
   ; mutable cps_apply_funs : Code.Var.t Stdlib.IntMap.t
   ; mutable curry_funs : Code.Var.t Stdlib.IntMap.t
@@ -57,7 +58,7 @@ val instr : Wasm_ast.instruction -> unit t
 
 val seq : unit t -> expression -> expression
 
-val expression_list : ('a -> expression) -> 'a list -> Wasm_ast.expression list t
+val expression_list : ('a -> 'b t) -> 'a list -> 'b list t
 
 module Arith : sig
   val const : int32 -> expression
@@ -198,3 +199,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 variable_type : Code.Var.t -> Wasm_ast.value_type option t

compiler/lib-wasm/gc_target.ml

@@ -281,11 +281,19 @@ module Type = struct
                     ])
             })
 
-  let env_type ~cps ~arity n =
+  let make_env_type env_type =
+    List.map
+      ~f:(fun typ ->
+        { W.mut = false
+        ; typ = W.Value (Option.value ~default:(W.Ref { nullable = false; typ = Eq }) typ)
+        })
+      env_type
+
+  let env_type ~cps ~arity ~env_type_id ~env_type =
     register_type
       (if cps
-       then Printf.sprintf "cps_env_%d_%d" arity n
-       else Printf.sprintf "env_%d_%d" arity n)
+       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
@@ -309,18 +317,11 @@ module Type = struct
                         ; typ = Value (Ref { nullable = false; typ = Type fun_ty' })
                         }
                       ])
-                @ List.init
-                    ~f:(fun _ ->
-                      { W.mut = false
-                      ; typ = W.Value (Ref { nullable = false; typ = Eq })
-                      })
-                    ~len:n)
+                @ make_env_type env_type)
           })
 
-  let rec_env_type ~function_count ~free_variable_count =
-    register_type
-      (Printf.sprintf "rec_env_%d_%d" function_count free_variable_count)
-      (fun () ->
+  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 () ->
         return
           { supertype = None
           ; final = true
@@ -331,24 +332,20 @@ module Type = struct
                      { W.mut = i < function_count
                      ; typ = W.Value (Ref { nullable = false; typ = Eq })
                      })
-                   ~len:(function_count + free_variable_count))
+                   ~len:function_count
+                @ make_env_type env_type)
           })
 
-  let rec_closure_type ~cps ~arity ~function_count ~free_variable_count =
+  let rec_closure_type ~cps ~arity ~function_count ~env_type_id ~env_type =
     register_type
       (if cps
-       then
-         Printf.sprintf
-           "cps_closure_rec_%d_%d_%d"
-           arity
-           function_count
-           free_variable_count
-       else Printf.sprintf "closure_rec_%d_%d_%d" arity function_count free_variable_count)
+       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 ~free_variable_count in
+        let* env_ty = rec_env_type ~function_count ~env_type_id ~env_type in
         return
           { supertype = Some cl_typ
           ; final = true
@@ -1099,11 +1096,19 @@ module Closure = struct
       in
       return (W.GlobalGet name)
     else
-      let free_variable_count = List.length free_variables in
+      let* env_type = expression_list variable_type free_variables in
+      let env_type_id =
+        try Hashtbl.find context.closure_types env_type
+        with Not_found ->
+          let id = Hashtbl.length context.closure_types in
+          Hashtbl.add context.closure_types env_type id;
+          id
+      in
+      info.id <- Some env_type_id;
       match info.Closure_conversion.functions with
       | [] -> assert false
       | [ _ ] ->
-          let* typ = Type.env_type ~cps ~arity free_variable_count in
+          let* typ = Type.env_type ~cps ~arity ~env_type_id ~env_type in
           let* l = expression_list load free_variables in
           return
             (W.StructNew
@@ -1122,7 +1127,7 @@ module Closure = struct
                  @ l ))
       | (g, _) :: _ as functions ->
           let function_count = List.length functions in
-          let* env_typ = Type.rec_env_type ~function_count ~free_variable_count in
+          let* env_typ = Type.rec_env_type ~function_count ~env_type_id ~env_type in
           let env =
             if Code.Var.equal f g
             then
@@ -1144,7 +1149,7 @@ module Closure = struct
               load env
           in
           let* typ =
-            Type.rec_closure_type ~cps ~arity ~function_count ~free_variable_count
+            Type.rec_closure_type ~cps ~arity ~function_count ~env_type_id ~env_type
           in
           let res =
             let* env = env in
@@ -1189,12 +1194,13 @@ module Closure = struct
       let* _ = add_var (Code.Var.fresh ()) in
       return ()
     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 offset = Memory.env_start arity in
       match info.Closure_conversion.functions with
       | [ _ ] ->
-          let* typ = Type.env_type ~cps ~arity free_variable_count in
+          let* typ = Type.env_type ~cps ~arity ~env_type_id ~env_type:[] in
           let* _ = add_var f in
           let env = Code.Var.fresh_n "env" in
           let* () =
@@ -1214,11 +1220,11 @@ module Closure = struct
       | functions ->
           let function_count = List.length functions in
           let* typ =
-            Type.rec_closure_type ~cps ~arity ~function_count ~free_variable_count
+            Type.rec_closure_type ~cps ~arity ~function_count ~env_type_id ~env_type:[]
           in
           let* _ = add_var f in
           let env = Code.Var.fresh_n "env" in
-          let* env_typ = Type.rec_env_type ~function_count ~free_variable_count in
+          let* env_typ = Type.rec_env_type ~function_count ~env_type_id ~env_type:[] in
           let* () =
             store
               ~typ:(W.Ref { nullable = false; typ = Type env_typ })