Delete StoreOpaque::traitobj_mut (bytecodealliance#11444)

* Delete `StoreOpaque::traitobj_mut`

This was a fundamentally unsound function because it is widening a
mutable borrow to encompass more than it originally contained. Removal
of its usage falls into a few buckets here:

* Many instances of `store.0.traitobj_mut()` were starting from a
  `&mut StoreInner<T>`, going to `&mut StoreOpaque`, then going back to
  `&mut dyn VMStore`. These are replaced with `store.0` as
  `&mut StoreInner<T>` directly coerces into `&mut dyn VMStore`.

* Some fiber-related helpers were updated to take any store-opaque-thing
  which encapsulates the ability to start/resume a borrow with a part of
  the store, but you only get that part of the store during the fiber
  itself. This means there's no widening necessary.

* Some methods previously taking `&mut StoreOpaque` are now
  appropriately widened to `&mut dyn VMStore`.

This all enables full deletion of this function which prevents
accidentally ever tripping over its unsound-ness.

* Fix configured build

* Another fix for a configured build

* Fix a warning

* Don't hold `VMStore` live over await points

* Review comments

* Move `AsStoreOpaque` to `store.rs` and use it in GC

* Fix configured build

* Switch lint annotation

Author: alexcrichton

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -1088,7 +1088,7 @@ impl Instance {
 
         impl<'a, T, V> Drop for Dropper<'a, T, V> {
             fn drop(&mut self) {
-                tls::set(self.store.0.traitobj_mut(), || {
+                tls::set(self.store.0, || {
                     // SAFETY: Here we drop the value without moving it for the
                     // first and only time -- per the contract for `Drop::drop`,
                     // this code won't run again, and the `value` field will no
@@ -1163,7 +1163,7 @@ impl Instance {
     /// can be made (in which case we trap with `Trap::AsyncDeadlock`).
     async fn poll_until<T, R>(
         self,
-        store: StoreContextMut<'_, T>,
+        mut store: StoreContextMut<'_, T>,
         mut future: Pin<&mut impl Future<Output = R>>,
     ) -> Result<R> {
         loop {
@@ -1312,19 +1312,23 @@ impl Instance {
                 // any work items and then loop again.
                 Either::Right(ready) => {
                     for item in ready {
-                        self.handle_work_item(store.0.traitobj_mut(), item).await?;
+                        self.handle_work_item(store.as_context_mut(), item).await?;
                     }
                 }
             }
         }
     }
 
     /// Handle the specified work item, possibly resuming a fiber if applicable.
-    async fn handle_work_item(self, store: &mut StoreOpaque, item: WorkItem) -> Result<()> {
+    async fn handle_work_item<T>(
+        self,
+        store: StoreContextMut<'_, T>,
+        item: WorkItem,
+    ) -> Result<()> {
         log::trace!("handle work item {item:?}");
         match item {
             WorkItem::PushFuture(future) => {
-                self.concurrent_state_mut(store)
+                self.concurrent_state_mut(store.0)
                     .futures
                     .get_mut()
                     .unwrap()
@@ -1333,10 +1337,10 @@ impl Instance {
                     .push(future.into_inner().unwrap());
             }
             WorkItem::ResumeFiber(fiber) => {
-                self.resume_fiber(store, fiber).await?;
+                self.resume_fiber(store.0, fiber).await?;
             }
             WorkItem::GuestCall(call) => {
-                let state = self.concurrent_state_mut(store);
+                let state = self.concurrent_state_mut(store.0);
                 if call.is_ready(state)? {
                     self.run_on_worker(store, WorkerItem::GuestCall(call))
                         .await?;
@@ -1362,7 +1366,7 @@ impl Instance {
                 }
             }
             WorkItem::Poll(params) => {
-                let state = self.concurrent_state_mut(store);
+                let state = self.concurrent_state_mut(store.0);
                 if state.get_mut(params.task)?.event.is_some()
                     || !state.get_mut(params.set)?.ready.is_empty()
                 {
@@ -1434,11 +1438,11 @@ impl Instance {
     }
 
     /// Execute the specified guest call on a worker fiber.
-    async fn run_on_worker(self, store: &mut StoreOpaque, item: WorkerItem) -> Result<()> {
-        let worker = if let Some(fiber) = self.concurrent_state_mut(store).worker.take() {
+    async fn run_on_worker<T>(self, store: StoreContextMut<'_, T>, item: WorkerItem) -> Result<()> {
+        let worker = if let Some(fiber) = self.concurrent_state_mut(store.0).worker.take() {
             fiber
         } else {
-            fiber::make_fiber(store.traitobj_mut(), move |store| {
+            fiber::make_fiber(store.0, move |store| {
                 loop {
                     match self.concurrent_state_mut(store).worker_item.take().unwrap() {
                         WorkerItem::GuestCall(call) => self.handle_guest_call(store, call)?,
@@ -1450,11 +1454,11 @@ impl Instance {
             })?
         };
 
-        let worker_item = &mut self.concurrent_state_mut(store).worker_item;
+        let worker_item = &mut self.concurrent_state_mut(store.0).worker_item;
         assert!(worker_item.is_none());
         *worker_item = Some(item);
 
-        self.resume_fiber(store, worker).await
+        self.resume_fiber(store.0, worker).await
     }
 
     /// Execute the specified guest call.
@@ -1531,7 +1535,7 @@ impl Instance {
         };
 
         let old_guest_task = if let Some(task) = task {
-            self.maybe_pop_call_context(store.store_opaque_mut(), task)?;
+            self.maybe_pop_call_context(store, task)?;
             self.concurrent_state_mut(store).guest_task
         } else {
             None
@@ -1545,7 +1549,7 @@ impl Instance {
 
         if let Some(task) = task {
             self.concurrent_state_mut(store).guest_task = old_guest_task;
-            self.maybe_push_call_context(store.store_opaque_mut(), task)?;
+            self.maybe_push_call_context(store, task)?;
         }
 
         Ok(())
@@ -2217,10 +2221,7 @@ impl Instance {
         // before committing to such an optimization.  And again, we'd need to
         // update the spec to allow that.
         let (status, waitable) = loop {
-            self.suspend(
-                store.0.traitobj_mut(),
-                SuspendReason::Waiting { set, task: caller },
-            )?;
+            self.suspend(store.0, SuspendReason::Waiting { set, task: caller })?;
 
             let state = self.concurrent_state_mut(store.0);
 
@@ -2414,7 +2415,7 @@ impl Instance {
             Poll::Ready(output) => {
                 // It finished immediately; lower the result and delete the
                 // task.
-                output.consume(store.0.traitobj_mut(), self)?;
+                output.consume(store.0, self)?;
                 log::trace!("delete host task {task:?} (already ready)");
                 self.concurrent_state_mut(store.0).delete(task)?;
                 None

crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs

@@ -680,11 +680,7 @@ impl<T> FutureReader<T> {
         // `self` should never be used again, but leave an invalid handle there just in case.
         let id = mem::replace(&mut self.id, TableId::new(0));
         self.instance
-            .host_drop_reader(
-                store.as_context_mut().0.traitobj_mut(),
-                id,
-                TransmitKind::Future,
-            )
+            .host_drop_reader(store.as_context_mut().0, id, TransmitKind::Future)
             .unwrap();
     }
 
@@ -1252,11 +1248,7 @@ impl<T> StreamReader<T> {
         // `self` should never be used again, but leave an invalid handle there just in case.
         let id = mem::replace(&mut self.id, TableId::new(0));
         self.instance
-            .host_drop_reader(
-                store.as_context_mut().0.traitobj_mut(),
-                id,
-                TransmitKind::Stream,
-            )
+            .host_drop_reader(store.as_context_mut().0, id, TransmitKind::Stream)
             .unwrap()
     }
 
@@ -2047,7 +2039,7 @@ impl Instance {
 
             WriteState::HostReady { accept, post_write } => {
                 accept(
-                    store.0.traitobj_mut(),
+                    store.0,
                     self,
                     Reader::Host {
                         accept: Box::new(|input, count| {
@@ -2863,7 +2855,7 @@ impl Instance {
                 }
 
                 let code = accept(
-                    store.0.traitobj_mut(),
+                    store.0,
                     self,
                     Reader::Guest {
                         options: &options,

crates/wasmtime/src/runtime/component/func/host.rs

@@ -344,7 +344,7 @@ where
             HostResult::Done(result) => result?,
             #[cfg(feature = "component-model-async")]
             HostResult::Future(future) => {
-                instance.poll_and_block(store.0.traitobj_mut(), future, caller_instance)?
+                instance.poll_and_block(store.0, future, caller_instance)?
             }
         };
 
@@ -826,8 +826,7 @@ where
             params_and_results,
             result_start,
         );
-        let result_vals =
-            instance.poll_and_block(store.0.traitobj_mut(), future, caller_instance)?;
+        let result_vals = instance.poll_and_block(store.0, future, caller_instance)?;
         let result_vals = &result_vals[result_start..];
 
         unsafe {

crates/wasmtime/src/runtime/component/linker.rs

@@ -837,7 +837,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
                         );
                         let mut future = std::pin::pin!(dtor(accessor, param));
                         std::future::poll_fn(|cx| {
-                            crate::component::concurrent::tls::set(store.0.traitobj_mut(), || {
+                            crate::component::concurrent::tls::set(store.0, || {
                                 future.as_mut().poll(cx)
                             })
                         })

crates/wasmtime/src/runtime/component/store.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "component-model-async")]
+use crate::runtime::vm::VMStore;
 use crate::runtime::vm::component::{ComponentInstance, OwnedComponentInstance};
 use crate::store::{StoreData, StoreId, StoreOpaque};
 #[cfg(feature = "component-model-async")]
@@ -32,7 +34,7 @@ impl ComponentStoreData {
     }
 
     #[cfg(feature = "component-model-async")]
-    pub(crate) fn drop_fibers_and_futures(store: &mut StoreOpaque) {
+    pub(crate) fn drop_fibers_and_futures(store: &mut dyn VMStore) {
         let mut fibers = Vec::new();
         let mut futures = Vec::new();
         for (_, instance) in store.store_data_mut().components.instances.iter_mut() {
@@ -50,7 +52,7 @@ impl ComponentStoreData {
             fiber.dispose(store);
         }
 
-        crate::component::concurrent::tls::set(store.traitobj_mut(), move || drop(futures));
+        crate::component::concurrent::tls::set(store, move || drop(futures));
     }
 }
 

crates/wasmtime/src/runtime/fiber.rs

@@ -1,7 +1,7 @@
 use crate::prelude::*;
-use crate::store::{Executor, StoreId, StoreInner, StoreOpaque};
+use crate::store::{AsStoreOpaque, Executor, StoreId, StoreOpaque};
 use crate::vm::mpk::{self, ProtectionMask};
-use crate::vm::{AlwaysMut, AsyncWasmCallState, VMStore};
+use crate::vm::{AlwaysMut, AsyncWasmCallState};
 use crate::{Engine, StoreContextMut};
 use anyhow::{Result, anyhow};
 use core::mem;
@@ -113,22 +113,6 @@ impl AsyncState {
     }
 }
 
-trait AsStoreOpaque {
-    fn as_store_opaque(&mut self) -> &mut StoreOpaque;
-}
-
-impl AsStoreOpaque for StoreOpaque {
-    fn as_store_opaque(&mut self) -> &mut StoreOpaque {
-        self
-    }
-}
-
-impl<T: 'static> AsStoreOpaque for StoreInner<T> {
-    fn as_store_opaque(&mut self) -> &mut StoreOpaque {
-        self
-    }
-}
-
 /// A helper structure used to block a fiber.
 ///
 /// This is acquired via either `StoreContextMut::with_blocking` or
@@ -785,15 +769,19 @@ fn resume_fiber<'a>(
 }
 
 /// Create a new `StoreFiber` which runs the specified closure.
-pub(crate) fn make_fiber<'a>(
-    store: &mut dyn VMStore,
-    fun: impl FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync + 'a,
-) -> Result<StoreFiber<'a>> {
-    let engine = store.engine().clone();
+pub(crate) fn make_fiber<'a, S>(
+    store: &mut S,
+    fun: impl FnOnce(&mut S) -> Result<()> + Send + Sync + 'a,
+) -> Result<StoreFiber<'a>>
+where
+    S: AsStoreOpaque + ?Sized + 'a,
+{
+    let opaque = store.as_store_opaque();
+    let engine = opaque.engine().clone();
     let executor = Executor::new(&engine);
-    let id = store.store_opaque().id();
-    let stack = store.store_opaque_mut().allocate_fiber_stack()?;
-    let track_pkey_context_switch = store.has_pkey();
+    let id = opaque.id();
+    let stack = opaque.allocate_fiber_stack()?;
+    let track_pkey_context_switch = opaque.has_pkey();
     let store = &raw mut *store;
     let fiber = Fiber::new(stack, move |result: WasmtimeResume, suspend| {
         let future_cx = match result {
@@ -812,17 +800,22 @@ pub(crate) fn make_fiber<'a>(
         // the fiber is running and `future_cx` and `suspend` are both in scope.
         // Note that these pointers are removed when this function returns as
         // that's when they fall out of scope.
-        let async_state = store_ref.store_opaque_mut().fiber_async_state_mut();
+        let async_state = store_ref.as_store_opaque().fiber_async_state_mut();
         assert!(async_state.current_suspend.is_none());
         assert!(async_state.current_future_cx.is_none());
         async_state.current_suspend = Some(NonNull::from(suspend));
         async_state.current_future_cx = Some(future_cx);
 
-        struct ResetCurrentPointersToNull<'a>(&'a mut dyn VMStore);
+        struct ResetCurrentPointersToNull<'a, S>(&'a mut S)
+        where
+            S: AsStoreOpaque + ?Sized;
 
-        impl Drop for ResetCurrentPointersToNull<'_> {
+        impl<S> Drop for ResetCurrentPointersToNull<'_, S>
+        where
+            S: AsStoreOpaque + ?Sized,
+        {
             fn drop(&mut self) {
-                let state = self.0.fiber_async_state_mut();
+                let state = self.0.as_store_opaque().fiber_async_state_mut();
 
                 // Double-check that the current suspension isn't null (it
                 // should be what's in this closure). Note though that we
@@ -861,23 +854,28 @@ pub(crate) fn make_fiber<'a>(
 
 /// Run the specified function on a newly-created fiber and `.await` its
 /// completion.
-pub(crate) async fn on_fiber<R: Send + Sync>(
-    store: &mut StoreOpaque,
-    func: impl FnOnce(&mut StoreOpaque) -> R + Send + Sync,
-) -> Result<R> {
-    let config = store.engine().config();
-    debug_assert!(store.async_support());
+pub(crate) async fn on_fiber<S, R>(
+    store: &mut S,
+    func: impl FnOnce(&mut S) -> R + Send + Sync,
+) -> Result<R>
+where
+    S: AsStoreOpaque + ?Sized,
+    R: Send + Sync,
+{
+    let opaque = store.as_store_opaque();
+    let config = opaque.engine().config();
+    debug_assert!(opaque.async_support());
     debug_assert!(config.async_stack_size > 0);
 
     let mut result = None;
-    let fiber = make_fiber(store.traitobj_mut(), |store| {
+    let fiber = make_fiber(store, |store| {
         result = Some(func(store));
         Ok(())
     })?;
 
     {
         let fiber = FiberFuture {
-            store,
+            store: store.as_store_opaque(),
             fiber: Some(fiber),
             #[cfg(feature = "component-model-async")]
             on_release: OnRelease::ReturnPending,