Optimize calling a WebAssembly function

crates/runtime/src/externref.rs

@@ -752,6 +752,7 @@ impl VMExternRefActivationsTable {
     /// // call has returned.
     /// drop(auto_reset_canary);
     /// ```
+    #[inline]
     pub fn set_stack_canary<'a>(&'a self, canary: &u8) -> impl Drop + 'a {
         let should_reset = if self.stack_canary.get().is_none() {
             let canary = canary as *const u8 as *mut u8;
@@ -775,6 +776,7 @@ impl VMExternRefActivationsTable {
         }
 
         impl Drop for AutoResetCanary<'_> {
+            #[inline]
             fn drop(&mut self) {
                 if self.should_reset {
                     debug_assert!(self.table.stack_canary.get().is_some());

crates/runtime/src/traphandlers.rs

@@ -4,8 +4,9 @@
 use crate::VMInterrupts;
 use backtrace::Backtrace;
 use std::any::Any;
-use std::cell::Cell;
+use std::cell::{Cell, UnsafeCell};
 use std::error::Error;
+use std::mem::MaybeUninit;
 use std::ptr;
 use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
 use std::sync::Once;
@@ -47,9 +48,10 @@ pub use sys::SignalHandler;
 /// function needs to be called at the end of the startup process, after other
 /// handlers have been installed. This function can thus be called multiple
 /// times, having no effect after the first call.
-pub fn init_traps() {
+pub fn init_traps() -> Result<(), Trap> {
     static INIT: Once = Once::new();
     INIT.call_once(|| unsafe { sys::platform_init() });
+    sys::lazy_per_thread_init()
 }
 
 /// Raises a user-defined trap immediately.
@@ -155,8 +157,6 @@ pub unsafe fn catch_traps<F>(trap_info: &impl TrapInfo, mut closure: F) -> Resul
 where
     F: FnMut(),
 {
-    sys::lazy_per_thread_init()?;
-
     return CallThreadState::new(trap_info).with(|cx| {
         RegisterSetjmp(
             cx.jmp_buf.as_ptr(),
@@ -191,7 +191,7 @@ pub fn out_of_gas() {
 /// Temporary state stored on the stack which is registered in the `tls` module
 /// below for calls into wasm.
 pub struct CallThreadState<'a> {
-    unwind: Cell<UnwindReason>,
+    unwind: UnsafeCell<MaybeUninit<UnwindReason>>,
     jmp_buf: Cell<*const u8>,
     handling_trap: Cell<bool>,
     trap_info: &'a (dyn TrapInfo + 'a),
@@ -232,17 +232,17 @@ pub unsafe trait TrapInfo {
 }
 
 enum UnwindReason {
-    None,
     Panic(Box<dyn Any + Send>),
     UserTrap(Box<dyn Error + Send + Sync>),
     LibTrap(Trap),
     JitTrap { backtrace: Backtrace, pc: usize },
 }
 
 impl<'a> CallThreadState<'a> {
+    #[inline]
     fn new(trap_info: &'a (dyn TrapInfo + 'a)) -> CallThreadState<'a> {
         CallThreadState {
-            unwind: Cell::new(UnwindReason::None),
+            unwind: UnsafeCell::new(MaybeUninit::uninit()),
             jmp_buf: Cell::new(ptr::null()),
             handling_trap: Cell::new(false),
             trap_info,
@@ -253,18 +253,13 @@ impl<'a> CallThreadState<'a> {
     fn with(self, closure: impl FnOnce(&CallThreadState) -> i32) -> Result<(), Trap> {
         let _reset = self.update_stack_limit()?;
         let ret = tls::set(&self, || closure(&self));
-        match self.unwind.replace(UnwindReason::None) {
-            UnwindReason::None => {
-                debug_assert_eq!(ret, 1);
-                Ok(())
-            }
-            UnwindReason::UserTrap(data) => {
-                debug_assert_eq!(ret, 0);
-                Err(Trap::User(data))
-            }
+        if ret != 0 {
+            return Ok(());
+        }
+        match unsafe { (*self.unwind.get()).as_ptr().read() } {
+            UnwindReason::UserTrap(data) => Err(Trap::User(data)),
             UnwindReason::LibTrap(trap) => Err(trap),
             UnwindReason::JitTrap { backtrace, pc } => {
-                debug_assert_eq!(ret, 0);
                 let interrupts = self.trap_info.interrupts();
                 let maybe_interrupted =
                     interrupts.stack_limit.load(SeqCst) == wasmtime_environ::INTERRUPTED;
@@ -274,10 +269,7 @@ impl<'a> CallThreadState<'a> {
                     maybe_interrupted,
                 })
             }
-            UnwindReason::Panic(panic) => {
-                debug_assert_eq!(ret, 0);
-                std::panic::resume_unwind(panic)
-            }
+            UnwindReason::Panic(panic) => std::panic::resume_unwind(panic),
         }
     }
 
@@ -310,6 +302,7 @@ impl<'a> CallThreadState<'a> {
     ///
     /// Note that this function must be called with `self` on the stack, not the
     /// heap/etc.
+    #[inline]
     fn update_stack_limit(&self) -> Result<impl Drop + '_, Trap> {
         // Determine the stack pointer where, after which, any wasm code will
         // immediately trap. This is checked on the entry to all wasm functions.
@@ -361,6 +354,7 @@ impl<'a> CallThreadState<'a> {
         struct Reset<'a>(bool, &'a AtomicUsize);
 
         impl Drop for Reset<'_> {
+            #[inline]
             fn drop(&mut self) {
                 if self.0 {
                     self.1.store(usize::max_value(), SeqCst);
@@ -372,8 +366,8 @@ impl<'a> CallThreadState<'a> {
     }
 
     fn unwind_with(&self, reason: UnwindReason) -> ! {
-        self.unwind.replace(reason);
         unsafe {
+            (*self.unwind.get()).as_mut_ptr().write(reason);
             Unwind(self.jmp_buf.get());
         }
     }
@@ -432,16 +426,21 @@ impl<'a> CallThreadState<'a> {
 
     fn capture_backtrace(&self, pc: *const u8) {
         let backtrace = Backtrace::new_unresolved();
-        self.unwind.replace(UnwindReason::JitTrap {
-            backtrace,
-            pc: pc as usize,
-        });
+        unsafe {
+            (*self.unwind.get())
+                .as_mut_ptr()
+                .write(UnwindReason::JitTrap {
+                    backtrace,
+                    pc: pc as usize,
+                });
+        }
     }
 }
 
 struct ResetCell<'a, T: Copy>(&'a Cell<T>, T);
 
 impl<T: Copy> Drop for ResetCell<'_, T> {
+    #[inline]
     fn drop(&mut self) {
         self.0.set(self.1);
     }
@@ -544,6 +543,7 @@ mod tls {
         struct Reset<'a, 'b>(&'a CallThreadState<'b>);
 
         impl Drop for Reset<'_, '_> {
+            #[inline]
             fn drop(&mut self) {
                 raw::replace(self.0.prev.replace(ptr::null()));
             }

crates/wasmtime/src/engine.rs

@@ -56,6 +56,7 @@ impl Engine {
     }
 
     /// Returns the configuration settings that this engine is using.
+    #[inline]
     pub fn config(&self) -> &Config {
         &self.inner.config
     }

crates/wasmtime/src/func.rs

@@ -804,7 +804,7 @@ impl Func {
     /// initiates a panic.
     pub fn call(&self, params: &[Val]) -> Result<Box<[Val]>> {
         assert!(
-            !self.store().async_support(),
+            !cfg!(feature = "async") || !self.store().async_support(),
             "must use `call_async` when async support is enabled on the config",
         );
         self._call(params)
@@ -926,6 +926,7 @@ impl Func {
     }
 
     /// Get a reference to this function's store.
+    #[inline]
     pub fn store(&self) -> &Store {
         &self.instance.store
     }
@@ -1414,6 +1415,7 @@ impl Caller<'_> {
     }
 
     /// Get a reference to the caller's store.
+    #[inline]
     pub fn store(&self) -> &Store {
         self.store
     }

crates/wasmtime/src/func/typed.rs

@@ -53,7 +53,7 @@ where
     /// connected to an asynchronous store.
     pub fn call(&self, params: Params) -> Result<Results, Trap> {
         assert!(
-            !self.func.store().async_support(),
+            !cfg!(feature = "async") || !self.func.store().async_support(),
             "must use `call_async` with async stores"
         );
         unsafe { self._call(params) }

crates/wasmtime/src/store.rs

@@ -136,7 +136,7 @@ impl Store {
         // once-per-thread. Platforms like Unix, however, only require this
         // once-per-program. In any case this is safe to call many times and
         // each one that's not relevant just won't do anything.
-        wasmtime_runtime::init_traps();
+        wasmtime_runtime::init_traps().expect("failed to initialize trap handling");
 
         Store {
             inner: Rc::new(StoreInner {
@@ -209,6 +209,7 @@ impl Store {
     }
 
     /// Returns the [`Engine`] that this store is associated with.
+    #[inline]
     pub fn engine(&self) -> &Engine {
         &self.inner.engine
     }
@@ -503,10 +504,12 @@ impl Store {
         }
     }
 
+    #[inline]
     pub(crate) fn externref_activations_table(&self) -> &VMExternRefActivationsTable {
         &self.inner.externref_activations_table
     }
 
+    #[inline]
     pub(crate) fn stack_map_registry(&self) -> &StackMapRegistry {
         &self.inner.stack_map_registry
     }
@@ -655,6 +658,7 @@ impl Store {
         });
     }
 
+    #[inline]
     pub(crate) fn async_support(&self) -> bool {
         self.inner.engine.config().async_support
     }
@@ -915,6 +919,7 @@ impl Store {
 }
 
 unsafe impl TrapInfo for Store {
+    #[inline]
     fn as_any(&self) -> &dyn Any {
         self
     }
@@ -930,6 +935,7 @@ unsafe impl TrapInfo for Store {
         false
     }
 
+    #[inline]
     fn max_wasm_stack(&self) -> usize {
         self.engine().config().max_wasm_stack
     }
@@ -956,6 +962,7 @@ unsafe impl TrapInfo for Store {
         }
     }
 
+    #[inline]
     fn interrupts(&self) -> &VMInterrupts {
         &self.inner.interrupts
     }