WIP use *const VMCallerCheckedAnyfunc as funcref

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,20 @@ 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);
+        pos.ins()
+            .trapz(anyfunc_ptr, ir::TrapCode::IndirectCallToNull);
+
+        // 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 +891,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 +908,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/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.as_ptr() },
+                    vmctx: unsafe { f.anyfunc.as_ref().vmctx },
                 });
             }
             (EntityIndex::Function(_), Some(_)) => {

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 {