Make the calls to wasm's memory.grow and memory.size indirect (#194)

* Make the calls to wasm's memory.grow and memory.size indirect

wasmtime-environ/src/func_environ.rs

@@ -8,9 +8,7 @@ use cranelift_codegen::ir;
 use cranelift_codegen::ir::condcodes::*;
 use cranelift_codegen::ir::immediates::{Offset32, Uimm64};
 use cranelift_codegen::ir::types::*;
-use cranelift_codegen::ir::{
+use cranelift_codegen::ir::{AbiParam, ArgumentPurpose, Function, InstBuilder, Signature};
-    AbiParam, ArgumentPurpose, ExtFuncData, FuncRef, Function, InstBuilder, Signature,
-};
 use cranelift_codegen::isa::TargetFrontendConfig;
 use cranelift_entity::EntityRef;
 use cranelift_wasm::{
@@ -50,6 +48,37 @@ pub fn get_imported_memory32_size_name() -> ir::ExternalName {
     ir::ExternalName::user(1, 3)
 }
 
+/// An index type for builtin functions.
+pub struct BuiltinFunctionIndex(u32);
+
+impl BuiltinFunctionIndex {
+    /// Returns an index for wasm's `memory.grow` builtin function.
+    pub const fn get_memory32_grow_index() -> Self {
+        Self(0)
+    }
+    /// Returns an index for wasm's imported `memory.grow` builtin function.
+    pub const fn get_imported_memory32_grow_index() -> Self {
+        Self(1)
+    }
+    /// Returns an index for wasm's `memory.size` builtin function.
+    pub const fn get_memory32_size_index() -> Self {
+        Self(2)
+    }
+    /// Returns an index for wasm's imported `memory.size` builtin function.
+    pub const fn get_imported_memory32_size_index() -> Self {
+        Self(3)
+    }
+    /// Returns the total number of builtin functions.
+    pub const fn builtin_functions_total_number() -> u32 {
+        4
+    }
+
+    /// Return the index as an u32 number.
+    pub const fn index(&self) -> u32 {
+        self.0
+    }
+}
+
 /// The `FuncEnvironment` implementation for use by the `ModuleEnvironment`.
 pub struct FuncEnvironment<'module_environment> {
     /// Target-specified configuration.
@@ -61,21 +90,13 @@ pub struct FuncEnvironment<'module_environment> {
     /// The Cranelift global holding the vmctx address.
     vmctx: Option<ir::GlobalValue>,
 
-    /// The external function declaration for implementing wasm's `memory.size`
+    /// The external function signature for implementing wasm's `memory.size`
     /// for locally-defined 32-bit memories.
-    memory32_size_extfunc: Option<FuncRef>,
+    memory32_size_sig: Option<ir::SigRef>,
 
-    /// The external function declaration for implementing wasm's `memory.size`
+    /// The external function signature for implementing wasm's `memory.grow`
-    /// for imported 32-bit memories.
-    imported_memory32_size_extfunc: Option<FuncRef>,
-
-    /// The external function declaration for implementing wasm's `memory.grow`
     /// for locally-defined memories.
-    memory_grow_extfunc: Option<FuncRef>,
+    memory_grow_sig: Option<ir::SigRef>,
-
-    /// The external function declaration for implementing wasm's `memory.grow`
-    /// for imported memories.
-    imported_memory_grow_extfunc: Option<FuncRef>,
 
     /// Offsets to struct fields accessed by JIT code.
     offsets: VMOffsets,
@@ -87,10 +108,8 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             target_config,
             module,
             vmctx: None,
-            memory32_size_extfunc: None,
+            memory32_size_sig: None,
-            imported_memory32_size_extfunc: None,
+            memory_grow_sig: None,
-            memory_grow_extfunc: None,
-            imported_memory_grow_extfunc: None,
             offsets: VMOffsets::new(target_config.pointer_bytes(), module),
         }
     }
@@ -107,106 +126,103 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         })
     }
 
-    fn get_memory_grow_sig(&self, func: &mut Function) -> ir::SigRef {
+    fn get_memory_grow_sig(&mut self, func: &mut Function) -> ir::SigRef {
-        func.import_signature(Signature {
+        let sig = self.memory_grow_sig.unwrap_or_else(|| {
-            params: vec![
+            func.import_signature(Signature {
-                AbiParam::special(self.pointer_type(), ArgumentPurpose::VMContext),
+                params: vec![
-                AbiParam::new(I32),
+                    AbiParam::special(self.pointer_type(), ArgumentPurpose::VMContext),
-                AbiParam::new(I32),
+                    AbiParam::new(I32),
-            ],
+                    AbiParam::new(I32),
-            returns: vec![AbiParam::new(I32)],
+                ],
-            call_conv: self.target_config.default_call_conv,
+                returns: vec![AbiParam::new(I32)],
-        })
+                call_conv: self.target_config.default_call_conv,
+            })
+        });
+        self.memory_grow_sig = Some(sig);
+        sig
     }
 
-    /// Return the memory.grow function to call for the given index, along with the
+    /// Return the memory.grow function signature to call for the given index, along with the
-    /// translated index value to pass to it.
+    /// translated index value to pass to it and its index in `VMBuiltinFunctionsArray`.
     fn get_memory_grow_func(
         &mut self,
         func: &mut Function,
         index: MemoryIndex,
-    ) -> (FuncRef, usize) {
+    ) -> (ir::SigRef, usize, BuiltinFunctionIndex) {
         if self.module.is_imported_memory(index) {
-            let extfunc = self.imported_memory_grow_extfunc.unwrap_or_else(|| {
-                let sig_ref = self.get_memory_grow_sig(func);
-                func.import_function(ExtFuncData {
-                    name: get_imported_memory32_grow_name(),
-                    signature: sig_ref,
-                    // We currently allocate all code segments independently, so nothing
-                    // is colocated.
-                    colocated: false,
-                })
-            });
-            self.imported_memory_grow_extfunc = Some(extfunc);
-            (extfunc, index.index())
-        } else {
-            let extfunc = self.memory_grow_extfunc.unwrap_or_else(|| {
-                let sig_ref = self.get_memory_grow_sig(func);
-                func.import_function(ExtFuncData {
-                    name: get_memory32_grow_name(),
-                    signature: sig_ref,
-                    // We currently allocate all code segments independently, so nothing
-                    // is colocated.
-                    colocated: false,
-                })
-            });
-            self.memory_grow_extfunc = Some(extfunc);
             (
-                extfunc,
+                self.get_memory_grow_sig(func),
+                index.index(),
+                BuiltinFunctionIndex::get_imported_memory32_grow_index(),
+            )
+        } else {
+            (
+                self.get_memory_grow_sig(func),
                 self.module.defined_memory_index(index).unwrap().index(),
+                BuiltinFunctionIndex::get_memory32_grow_index(),
             )
         }
     }
 
-    fn get_memory32_size_sig(&self, func: &mut Function) -> ir::SigRef {
+    fn get_memory32_size_sig(&mut self, func: &mut Function) -> ir::SigRef {
-        func.import_signature(Signature {
+        let sig = self.memory32_size_sig.unwrap_or_else(|| {
-            params: vec![
+            func.import_signature(Signature {
-                AbiParam::special(self.pointer_type(), ArgumentPurpose::VMContext),
+                params: vec![
-                AbiParam::new(I32),
+                    AbiParam::special(self.pointer_type(), ArgumentPurpose::VMContext),
-            ],
+                    AbiParam::new(I32),
-            returns: vec![AbiParam::new(I32)],
+                ],
-            call_conv: self.target_config.default_call_conv,
+                returns: vec![AbiParam::new(I32)],
-        })
+                call_conv: self.target_config.default_call_conv,
+            })
+        });
+        self.memory32_size_sig = Some(sig);
+        sig
     }
 
-    /// Return the memory.size function to call for the given index, along with the
+    /// Return the memory.size function signature to call for the given index, along with the
-    /// translated index value to pass to it.
+    /// translated index value to pass to it and its index in `VMBuiltinFunctionsArray`.
     fn get_memory_size_func(
         &mut self,
         func: &mut Function,
         index: MemoryIndex,
-    ) -> (FuncRef, usize) {
+    ) -> (ir::SigRef, usize, BuiltinFunctionIndex) {
         if self.module.is_imported_memory(index) {
-            let extfunc = self.imported_memory32_size_extfunc.unwrap_or_else(|| {
-                let sig_ref = self.get_memory32_size_sig(func);
-                func.import_function(ExtFuncData {
-                    name: get_imported_memory32_size_name(),
-                    signature: sig_ref,
-                    // We currently allocate all code segments independently, so nothing
-                    // is colocated.
-                    colocated: false,
-                })
-            });
-            self.imported_memory32_size_extfunc = Some(extfunc);
-            (extfunc, index.index())
-        } else {
-            let extfunc = self.memory32_size_extfunc.unwrap_or_else(|| {
-                let sig_ref = self.get_memory32_size_sig(func);
-                func.import_function(ExtFuncData {
-                    name: get_memory32_size_name(),
-                    signature: sig_ref,
-                    // We currently allocate all code segments independently, so nothing
-                    // is colocated.
-                    colocated: false,
-                })
-            });
-            self.memory32_size_extfunc = Some(extfunc);
             (
-                extfunc,
+                self.get_memory32_size_sig(func),
+                index.index(),
+                BuiltinFunctionIndex::get_imported_memory32_size_index(),
+            )
+        } else {
+            (
+                self.get_memory32_size_sig(func),
                 self.module.defined_memory_index(index).unwrap().index(),
+                BuiltinFunctionIndex::get_memory32_size_index(),
             )
         }
     }
+
+    /// Translates load of builtin function and returns a pair of values `vmctx`
+    /// and address of the loaded function.
+    fn translate_load_builtin_function_address(
+        &mut self,
+        pos: &mut FuncCursor<'_>,
+        callee_func_idx: BuiltinFunctionIndex,
+    ) -> (ir::Value, ir::Value) {
+        // We use an indirect call so that we don't have to patch the code at runtime.
+        let pointer_type = self.pointer_type();
+        let vmctx = self.vmctx(&mut pos.func);
+        let base = pos.ins().global_value(pointer_type, vmctx);
+
+        let mut mem_flags = ir::MemFlags::trusted();
+        mem_flags.set_readonly();
+
+        // Load the callee address.
+        let body_offset =
+            i32::try_from(self.offsets.vmctx_builtin_function(callee_func_idx)).unwrap();
+        let func_addr = pos.ins().load(pointer_type, mem_flags, base, body_offset);
+
+        (base, func_addr)
+    }
 }
 
 #[cfg(feature = "lightbeam")]
@@ -662,12 +678,12 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         _heap: ir::Heap,
         val: ir::Value,
     ) -> WasmResult<ir::Value> {
-        let (memory_grow_func, index_arg) = self.get_memory_grow_func(&mut pos.func, index);
+        let (func_sig, index_arg, func_idx) = self.get_memory_grow_func(&mut pos.func, index);
         let memory_index = pos.ins().iconst(I32, index_arg as i64);
-        let vmctx = pos.func.special_param(ArgumentPurpose::VMContext).unwrap();
+        let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
         let call_inst = pos
             .ins()
-            .call(memory_grow_func, &[vmctx, val, memory_index]);
+            .call_indirect(func_sig, func_addr, &[vmctx, val, memory_index]);
         Ok(*pos.func.dfg.inst_results(call_inst).first().unwrap())
     }
 
@@ -677,10 +693,12 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         index: MemoryIndex,
         _heap: ir::Heap,
     ) -> WasmResult<ir::Value> {
-        let (memory_size_func, index_arg) = self.get_memory_size_func(&mut pos.func, index);
+        let (func_sig, index_arg, func_idx) = self.get_memory_size_func(&mut pos.func, index);
         let memory_index = pos.ins().iconst(I32, index_arg as i64);
-        let vmctx = pos.func.special_param(ArgumentPurpose::VMContext).unwrap();
+        let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
-        let call_inst = pos.ins().call(memory_size_func, &[vmctx, memory_index]);
+        let call_inst = pos
+            .ins()
+            .call_indirect(func_sig, func_addr, &[vmctx, memory_index]);
         Ok(*pos.func.dfg.inst_results(call_inst).first().unwrap())
     }
 }

wasmtime-environ/src/lib.rs

@@ -54,6 +54,7 @@ pub use crate::compilation::{
     Compilation, CompileError, Compiler, Relocation, RelocationTarget, Relocations,
 };
 pub use crate::cranelift::Cranelift;
+pub use crate::func_environ::BuiltinFunctionIndex;
 #[cfg(feature = "lightbeam")]
 pub use crate::lightbeam::Lightbeam;
 pub use crate::module::{

wasmtime-environ/src/vmoffsets.rs

@@ -2,6 +2,7 @@
 //! module.
 
 use crate::module::Module;
+use crate::BuiltinFunctionIndex;
 use core::convert::TryFrom;
 use cranelift_codegen::ir;
 use cranelift_wasm::{
@@ -332,9 +333,8 @@ impl VMOffsets {
             .unwrap()
     }
 
-    /// Return the size of the `VMContext` allocation.
+    /// The offset of the builtin functions array.
-    #[allow(dead_code)]
+    pub fn vmctx_builtin_functions_begin(&self) -> u32 {
-    pub fn size_of_vmctx(&self) -> u32 {
         self.vmctx_globals_begin()
             .checked_add(
                 self.num_defined_globals
@@ -344,6 +344,17 @@ impl VMOffsets {
             .unwrap()
     }
 
+    /// Return the size of the `VMContext` allocation.
+    pub fn size_of_vmctx(&self) -> u32 {
+        self.vmctx_builtin_functions_begin()
+            .checked_add(
+                BuiltinFunctionIndex::builtin_functions_total_number()
+                    .checked_mul(u32::from(self.pointer_size))
+                    .unwrap(),
+            )
+            .unwrap()
+    }
+
     /// Return the offset to `VMSharedSignatureId` index `index`.
     pub fn vmctx_vmshared_signature_id(&self, index: SignatureIndex) -> u32 {
         assert!(index.as_u32() < self.num_signature_ids);
@@ -517,6 +528,18 @@ impl VMOffsets {
             .checked_add(u32::from(self.vmglobal_import_from()))
             .unwrap()
     }
+
+    /// Return the offset to builtin function in `VMBuiltinFunctionsArray` index `index`.
+    pub fn vmctx_builtin_function(&self, index: BuiltinFunctionIndex) -> u32 {
+        self.vmctx_builtin_functions_begin()
+            .checked_add(
+                index
+                    .index()
+                    .checked_mul(u32::from(self.pointer_size))
+                    .unwrap(),
+            )
+            .unwrap()
+    }
 }
 
 /// Target specific type for shared signature index.

wasmtime-runtime/src/instance.rs

@@ -11,9 +11,9 @@ use crate::signalhandlers::{wasmtime_init_eager, wasmtime_init_finish};
 use crate::table::Table;
 use crate::traphandlers::wasmtime_call;
 use crate::vmcontext::{
-    VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMFunctionImport, VMGlobalDefinition,
+    VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMFunctionImport,
-    VMGlobalImport, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex, VMTableDefinition,
+    VMGlobalDefinition, VMGlobalImport, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex,
-    VMTableImport,
+    VMTableDefinition, VMTableImport,
 };
 use core::any::Any;
 use core::borrow::Borrow;
@@ -360,6 +360,15 @@ impl Instance {
         }
     }
 
+    /// Return a pointer to the `VMBuiltinFunctionsArray`.
+    fn builtin_functions_ptr(&mut self) -> *mut VMBuiltinFunctionsArray {
+        unsafe {
+            (&mut self.vmctx as *mut VMContext as *mut u8)
+                .add(usize::try_from(self.offsets.vmctx_builtin_functions_begin()).unwrap())
+                as *mut VMBuiltinFunctionsArray
+        }
+    }
+
     /// Return a reference to the vmctx used by compiled wasm code.
     pub fn vmctx(&self) -> &VMContext {
         &self.vmctx
@@ -723,6 +732,10 @@ impl InstanceHandle {
                 instance.globals_ptr() as *mut VMGlobalDefinition,
                 vmctx_globals.len(),
             );
+            ptr::write(
+                instance.builtin_functions_ptr() as *mut VMBuiltinFunctionsArray,
+                VMBuiltinFunctionsArray::initialized(),
+            );
         }
 
         // Check initializer bounds before initializing anything.

wasmtime-runtime/src/vmcontext.rs

@@ -4,6 +4,7 @@
 use crate::instance::Instance;
 use core::any::Any;
 use core::{ptr, u32};
+use wasmtime_environ::BuiltinFunctionIndex;
 
 /// An imported function.
 #[derive(Debug, Copy, Clone)]
@@ -472,6 +473,33 @@ impl Default for VMCallerCheckedAnyfunc {
     }
 }
 
+/// An array that stores addresses of builtin functions. We translate code
+/// to use indirect calls. This way, we don't have to patch the code.
+#[repr(C)]
+pub struct VMBuiltinFunctionsArray {
+    ptrs: [usize; Self::len()],
+}
+
+impl VMBuiltinFunctionsArray {
+    pub const fn len() -> usize {
+        BuiltinFunctionIndex::builtin_functions_total_number() as usize
+    }
+
+    pub fn initialized() -> Self {
+        use crate::libcalls::*;
+        let mut ptrs = [0; Self::len()];
+        ptrs[BuiltinFunctionIndex::get_memory32_grow_index().index() as usize] =
+            wasmtime_memory32_grow as usize;
+        ptrs[BuiltinFunctionIndex::get_imported_memory32_grow_index().index() as usize] =
+            wasmtime_imported_memory32_grow as usize;
+        ptrs[BuiltinFunctionIndex::get_memory32_size_index().index() as usize] =
+            wasmtime_memory32_size as usize;
+        ptrs[BuiltinFunctionIndex::get_imported_memory32_size_index().index() as usize] =
+            wasmtime_imported_memory32_size as usize;
+        Self { ptrs }
+    }
+}
+
 /// The VM "context", which is pointed to by the `vmctx` arg in Cranelift.
 /// This has information about globals, memories, tables, and other runtime
 /// state associated with the current instance.