wasmtime: Implement funcref with NonNull<VMCallerCheckedAnyfunc>

This should be more efficient than using a `VMExternRef` that points at a
`VMCallerCheckedAnyfunc` because it gets rid of an indirection, dynamic
allocation, and some reference counting.

Note that the null function reference is *NOT* a null pointer; it is a
`VMCallerCheckedAnyfunc` that has a null `func_ptr` member.

cranelift/wasm/src/sections_translator.rs

@@ -97,6 +97,7 @@ pub fn parse_import_section<'data>(
             ImportSectionEntryType::Global(ref ty) => {
                 environ.declare_global_import(
                     Global {
+                        wasm_ty: ty.content_type,
                         ty: type_to_type(ty.content_type, environ).unwrap(),
                         mutability: ty.mutable,
                         initializer: GlobalInit::Import,
@@ -227,6 +228,7 @@ pub fn parse_global_section(
             }
         };
         let global = Global {
+            wasm_ty: content_type,
             ty: type_to_type(content_type, environ).unwrap(),
             mutability: mutable,
             initializer,

cranelift/wasm/src/translation_utils.rs

@@ -67,10 +67,18 @@ entity_impl!(DataIndex);
 pub struct ElemIndex(u32);
 entity_impl!(ElemIndex);
 
-/// WebAssembly global.
+/// A WebAssembly global.
+///
+/// Note that we record both the original Wasm type and the Cranelift IR type
+/// used to represent it. This is because multiple different kinds of Wasm types
+/// might be represented with the same Cranelift IR type. For example, both a
+/// Wasm `i64` and a `funcref` might be represented with a Cranelift `i64` on
+/// 64-bit architectures, and when GC is not required for func refs.
 #[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
 pub struct Global {
-    /// The type of the value stored in the global.
+    /// The Wasm type of the value stored in the global.
+    pub wasm_ty: crate::WasmType,
+    /// The Cranelift IR type of the value stored in the global.
     pub ty: ir::Type,
     /// A flag indicating whether the value may change at runtime.
     pub mutability: bool,

crates/environ/src/func_environ.rs

@@ -10,8 +10,8 @@ use cranelift_codegen::ir::{AbiParam, ArgumentPurpose, Function, InstBuilder, Si
 use cranelift_codegen::isa::{self, TargetFrontendConfig};
 use cranelift_entity::EntityRef;
 use cranelift_wasm::{
-    self, FuncIndex, GlobalIndex, GlobalVariable, MemoryIndex, SignatureIndex, TableIndex,
-    TargetEnvironment, WasmError, WasmResult,
+    self, FuncIndex, GlobalIndex, GlobalVariable, MemoryIndex, SignatureIndex, TableElementType,
+    TableIndex, TargetEnvironment, WasmError, WasmResult,
 };
 #[cfg(feature = "lightbeam")]
 use cranelift_wasm::{DefinedFuncIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex};
@@ -68,6 +68,10 @@ macro_rules! declare_builtin_functions {
                 AbiParam::new(self.reference_type)
             }
 
+            fn pointer(&self) -> AbiParam {
+                AbiParam::new(self.pointer_type)
+            }
+
             fn i32(&self) -> AbiParam {
                 AbiParam::new(I32)
             }
@@ -161,10 +165,10 @@ declare_builtin_functions! {
     memory_init(vmctx, i32, i32, i32, i32, i32) -> ();
     /// Returns an index for wasm's `data.drop` instruction.
     data_drop(vmctx, i32) -> ();
-    /// Returns an index for Wasm's `table.grow` instruction.
-    table_grow(vmctx, i32, i32, reference) -> (i32);
-    /// Returns an index for Wasm's `ref.func` instruction.
-    ref_func(vmctx, i32) -> (reference);
+    /// Returns an index for Wasm's `table.grow` instruction for `funcref`s.
+    table_grow_funcref(vmctx, i32, i32, pointer) -> (i32);
+    /// Returns an index for Wasm's `table.grow` instruction for `externref`s.
+    table_grow_externref(vmctx, i32, i32, reference) -> (i32);
 }
 
 impl BuiltinFunctionIndex {
@@ -586,9 +590,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         });
 
         let element_size = match self.module.table_plans[index].style {
-            TableStyle::CallerChecksSignature => {
-                u64::from(self.offsets.size_of_vmcaller_checked_anyfunc())
-            }
+            TableStyle::CallerChecksSignature => self.pointer_type().bytes() as _,
         };
 
         Ok(func.create_table(ir::TableData {
@@ -607,12 +609,25 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         delta: ir::Value,
         init_value: ir::Value,
     ) -> WasmResult<ir::Value> {
-        let table_index_arg = pos.ins().iconst(I32, table_index.as_u32() as i64);
+        let (func_idx, func_sig) = match self.module.table_plans[table_index].table.ty {
+            TableElementType::Func => (
+                BuiltinFunctionIndex::table_grow_funcref(),
+                self.builtin_function_signatures
+                    .table_grow_funcref(&mut pos.func),
+            ),
+            TableElementType::Val(ty) => {
+                debug_assert_eq!(ty, self.reference_type());
+                (
+                    BuiltinFunctionIndex::table_grow_externref(),
+                    self.builtin_function_signatures
+                        .table_grow_externref(&mut pos.func),
+                )
+            }
+        };
 
-        let func_idx = BuiltinFunctionIndex::table_grow();
         let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
 
-        let func_sig = self.builtin_function_signatures.table_grow(&mut pos.func);
+        let table_index_arg = pos.ins().iconst(I32, table_index.as_u32() as i64);
         let call_inst = pos.ins().call_indirect(
             func_sig,
             func_addr,
@@ -663,17 +678,10 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         mut pos: cranelift_codegen::cursor::FuncCursor<'_>,
         func_index: FuncIndex,
     ) -> WasmResult<ir::Value> {
-        let func_sig = self.builtin_function_signatures.ref_func(&mut pos.func);
-        let ref_func_index = BuiltinFunctionIndex::ref_func();
-        let func_index_arg = pos.ins().iconst(I32, func_index.as_u32() as i64);
-
-        let (vmctx, func_addr) =
-            self.translate_load_builtin_function_address(&mut pos, ref_func_index);
-        let call_inst = pos
-            .ins()
-            .call_indirect(func_sig, func_addr, &[vmctx, func_index_arg]);
-
-        Ok(pos.func.dfg.first_result(call_inst))
+        let vmctx = self.vmctx(&mut pos.func);
+        let vmctx = pos.ins().global_value(self.pointer_type(), vmctx);
+        let offset = self.offsets.vmctx_anyfunc(func_index);
+        Ok(pos.ins().iadd_imm(vmctx, i64::from(offset)))
     }
 
     fn translate_custom_global_get(
@@ -843,12 +851,21 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
 
         let table_entry_addr = pos.ins().table_addr(pointer_type, table, callee, 0);
 
-        // Dereference table_entry_addr to get the function address.
+        // Dereference the table entry to get the pointer to the
+        // `VMCallerCheckedAnyfunc`.
+        let anyfunc_ptr =
+            pos.ins()
+                .load(pointer_type, ir::MemFlags::trusted(), table_entry_addr, 0);
+        // Note: `anyfunc_ptr` is always non-null, even for null func refs! In
+        // that case, the pointed-to `VMCallerCheckedAnyfunc`'s `func_ptr`
+        // member is null, which is checked below.
+
+        // Dereference anyfunc pointer to get the function address.
         let mem_flags = ir::MemFlags::trusted();
         let func_addr = pos.ins().load(
             pointer_type,
             mem_flags,
-            table_entry_addr,
+            anyfunc_ptr,
             i32::from(self.offsets.vmcaller_checked_anyfunc_func_ptr()),
         );
 
@@ -875,7 +892,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                 let callee_sig_id = pos.ins().load(
                     sig_id_type,
                     mem_flags,
-                    table_entry_addr,
+                    anyfunc_ptr,
                     i32::from(self.offsets.vmcaller_checked_anyfunc_type_index()),
                 );
 
@@ -892,7 +909,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let vmctx = pos.ins().load(
             pointer_type,
             mem_flags,
-            table_entry_addr,
+            anyfunc_ptr,
             i32::from(self.offsets.vmcaller_checked_anyfunc_vmctx()),
         );
         real_call_args.push(vmctx);

crates/environ/src/vmoffsets.rs

@@ -15,6 +15,7 @@
 //      tables: [VMTableDefinition; module.num_defined_tables],
 //      memories: [VMMemoryDefinition; module.num_defined_memories],
 //      globals: [VMGlobalDefinition; module.num_defined_globals],
+//      anyfuncs: [VMCallerCheckedAnyfunc; module.num_imported_functions + module.num_defined_functions],
 //      builtins: VMBuiltinFunctionsArray,
 // }
 
@@ -62,6 +63,8 @@ pub struct VMOffsets {
     pub num_imported_memories: u32,
     /// The number of imported globals in the module.
     pub num_imported_globals: u32,
+    /// The number of defined functions in the module.
+    pub num_defined_functions: u32,
     /// The number of defined tables in the module.
     pub num_defined_tables: u32,
     /// The number of defined memories in the module.
@@ -80,6 +83,7 @@ impl VMOffsets {
             num_imported_tables: cast_to_u32(module.num_imported_tables),
             num_imported_memories: cast_to_u32(module.num_imported_memories),
             num_imported_globals: cast_to_u32(module.num_imported_globals),
+            num_defined_functions: cast_to_u32(module.functions.len()),
             num_defined_tables: cast_to_u32(module.table_plans.len()),
             num_defined_memories: cast_to_u32(module.memory_plans.len()),
             num_defined_globals: cast_to_u32(module.globals.len()),
@@ -390,8 +394,8 @@ impl VMOffsets {
         align(offset, 16)
     }
 
-    /// The offset of the builtin functions array.
-    pub fn vmctx_builtin_functions_begin(&self) -> u32 {
+    /// The offset of the `anyfuncs` array.
+    pub fn vmctx_anyfuncs_begin(&self) -> u32 {
         self.vmctx_globals_begin()
             .checked_add(
                 self.num_defined_globals
@@ -401,6 +405,19 @@ impl VMOffsets {
             .unwrap()
     }
 
+    /// The offset of the builtin functions array.
+    pub fn vmctx_builtin_functions_begin(&self) -> u32 {
+        self.vmctx_anyfuncs_begin()
+            .checked_add(
+                self.num_imported_functions
+                    .checked_add(self.num_defined_functions)
+                    .unwrap()
+                    .checked_mul(u32::from(self.size_of_vmcaller_checked_anyfunc()))
+                    .unwrap(),
+            )
+            .unwrap()
+    }
+
     /// Return the size of the `VMContext` allocation.
     pub fn size_of_vmctx(&self) -> u32 {
         self.vmctx_builtin_functions_begin()
@@ -516,6 +533,19 @@ impl VMOffsets {
             .unwrap()
     }
 
+    /// Return the offset to the `VMCallerCheckedAnyfunc` for the given function
+    /// index (either imported or defined).
+    pub fn vmctx_anyfunc(&self, index: FuncIndex) -> u32 {
+        self.vmctx_anyfuncs_begin()
+            .checked_add(
+                index
+                    .as_u32()
+                    .checked_mul(u32::from(self.size_of_vmcaller_checked_anyfunc()))
+                    .unwrap(),
+            )
+            .unwrap()
+    }
+
     /// Return the offset to the `body` field in `*const VMFunctionBody` index `index`.
     pub fn vmctx_vmfunction_import_body(&self, index: FuncIndex) -> u32 {
         self.vmctx_vmfunction_import(index)

crates/fuzzing/src/oracles/dummy.rs

@@ -13,7 +13,7 @@ pub fn dummy_imports<'module>(
     import_tys
         .map(|imp| {
             Ok(match imp.ty() {
-                ExternType::Func(func_ty) => Extern::Func(dummy_func(&store, func_ty)),
+                ExternType::Func(func_ty) => Extern::Func(Some(dummy_func(&store, func_ty))),
                 ExternType::Global(global_ty) => Extern::Global(dummy_global(&store, global_ty)?),
                 ExternType::Table(table_ty) => Extern::Table(dummy_table(&store, table_ty)?),
                 ExternType::Memory(mem_ty) => Extern::Memory(dummy_memory(&store, mem_ty)),

crates/jit/src/imports.rs

@@ -32,7 +32,9 @@ pub fn resolve_imports(
         match (import, &export) {
             (EntityIndex::Function(func_index), Some(Export::Function(f))) => {
                 let import_signature = module.local.native_func_signature(*func_index);
-                let signature = signatures.lookup_native(f.signature).unwrap();
+                let signature = signatures
+                    .lookup_native(unsafe { f.anyfunc.as_ref().type_index })
+                    .unwrap();
                 if signature != *import_signature {
                     // TODO: If the difference is in the calling convention,
                     // we could emit a wrapper function to fix it up.
@@ -43,8 +45,8 @@ pub fn resolve_imports(
                     )));
                 }
                 function_imports.push(VMFunctionImport {
-                    body: f.address,
-                    vmctx: f.vmctx,
+                    body: unsafe { f.anyfunc.as_ref().func_ptr },
+                    vmctx: unsafe { f.anyfunc.as_ref().vmctx },
                 });
             }
             (EntityIndex::Function(_), Some(_)) => {
@@ -169,16 +171,20 @@ fn is_global_compatible(exported: &Global, imported: &Global) -> bool {
     }
 
     let Global {
+        wasm_ty: exported_wasm_ty,
         ty: exported_ty,
         mutability: exported_mutability,
         initializer: _exported_initializer,
     } = exported;
     let Global {
+        wasm_ty: imported_wasm_ty,
         ty: imported_ty,
         mutability: imported_mutability,
         initializer: _imported_initializer,
     } = imported;
-    exported_ty == imported_ty && imported_mutability == exported_mutability
+    exported_wasm_ty == imported_wasm_ty
+        && exported_ty == imported_ty
+        && imported_mutability == exported_mutability
 }
 
 fn is_table_element_type_compatible(

crates/runtime/src/export.rs

@@ -1,7 +1,7 @@
 use crate::vmcontext::{
-    VMContext, VMFunctionBody, VMGlobalDefinition, VMMemoryDefinition, VMSharedSignatureIndex,
-    VMTableDefinition,
+    VMCallerCheckedAnyfunc, VMContext, VMGlobalDefinition, VMMemoryDefinition, VMTableDefinition,
 };
+use std::ptr::NonNull;
 use wasmtime_environ::wasm::Global;
 use wasmtime_environ::{MemoryPlan, TablePlan};
 
@@ -24,14 +24,8 @@ pub enum Export {
 /// A function export value.
 #[derive(Debug, Clone)]
 pub struct ExportFunction {
-    /// The address of the native-code function.
-    pub address: *const VMFunctionBody,
-    /// Pointer to the containing `VMContext`.
-    pub vmctx: *mut VMContext,
-    /// The function signature declaration, used for compatibilty checking.
-    ///
-    /// Note that this indexes within the module associated with `vmctx`.
-    pub signature: VMSharedSignatureIndex,
+    /// The `VMCallerCheckedAnyfunc` for this exported function.
+    pub anyfunc: NonNull<VMCallerCheckedAnyfunc>,
 }
 
 impl From<ExportFunction> for Export {