Refactor the wasmtime::runtime::vm::Store trait (#9277)

* Rename the `wasmtime::runtime::vm::Store` trait to `VMStore`

Purely mechanical.

* Add `StoreOpaque` accessors to the `VMStore` trait

* Make `dyn VMStore` deref to `StoreOpaque`

* Remove the `VMStore::vmruntime_limits` method

* Remove the `VMStore::epoch_ptr` method

* Rename `VMStore::gc_store` to `unwrap_gc_store_mut`

Will make removing the trait method easier.

* Remove `VMStore`'s `GcStore`-related methods

* cargo fmt

* Fix build with shared memory disabled

crates/wasmtime/src/runtime/store.rs

@@ -102,7 +102,6 @@ use core::num::NonZeroU64;
 use core::ops::{Deref, DerefMut};
 use core::pin::Pin;
 use core::ptr;
-use core::sync::atomic::AtomicU64;
 use core::task::{Context, Poll};
 
 mod context;
@@ -593,8 +592,8 @@ impl<T> Store<T> {
             // throughout Wasmtime.
             unsafe {
                 let traitobj = mem::transmute::<
-                    *mut (dyn crate::runtime::vm::Store + '_),
-                    *mut (dyn crate::runtime::vm::Store + 'static),
+                    *mut (dyn crate::runtime::vm::VMStore + '_),
+                    *mut (dyn crate::runtime::vm::VMStore + 'static),
                 >(&mut *inner);
                 instance.set_store(traitobj);
             }
@@ -1883,7 +1882,7 @@ impl StoreOpaque {
         self.default_caller.vmctx()
     }
 
-    pub fn traitobj(&self) -> *mut dyn crate::runtime::vm::Store {
+    pub fn traitobj(&self) -> *mut dyn crate::runtime::vm::VMStore {
         self.default_caller.store()
     }
 
@@ -2423,17 +2422,13 @@ impl AsyncCx {
     }
 }
 
-unsafe impl<T> crate::runtime::vm::Store for StoreInner<T> {
-    fn vmruntime_limits(&self) -> *mut VMRuntimeLimits {
-        <StoreOpaque>::vmruntime_limits(self)
-    }
-
-    fn epoch_ptr(&self) -> *const AtomicU64 {
-        self.engine.epoch_counter() as *const _
+unsafe impl<T> crate::runtime::vm::VMStore for StoreInner<T> {
+    fn store_opaque(&self) -> &StoreOpaque {
+        &self.inner
     }
 
-    fn maybe_gc_store(&mut self) -> Option<&mut GcStore> {
-        self.gc_store.as_mut()
+    fn store_opaque_mut(&mut self) -> &mut StoreOpaque {
+        &mut self.inner
     }
 
     fn memory_growing(

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

@@ -32,7 +32,7 @@ impl<'a, T> StoreContextMut<'a, T> {
     /// Unfortunately there's not a ton of debug asserts we can add here, so we
     /// rely on testing to largely help show that this is correctly used.
     pub(crate) unsafe fn from_raw(
-        store: *mut dyn crate::runtime::vm::Store,
+        store: *mut dyn crate::runtime::vm::VMStore,
     ) -> StoreContextMut<'a, T> {
         StoreContextMut(&mut *(store as *mut StoreInner<T>))
     }

crates/wasmtime/src/runtime/vm.rs

@@ -4,11 +4,14 @@
 #![warn(clippy::cast_sign_loss)]
 
 use crate::prelude::*;
+use crate::store::StoreOpaque;
 use alloc::sync::Arc;
 use core::fmt;
 use core::mem;
+use core::ops::Deref;
+use core::ops::DerefMut;
 use core::ptr::NonNull;
-use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicUsize, Ordering};
 use wasmtime_environ::{
     DefinedFuncIndex, DefinedMemoryIndex, HostPtr, ModuleInternedTypeIndex, VMOffsets,
     VMSharedTypeIndex,
@@ -92,27 +95,12 @@ pub use crate::runtime::vm::cow::{MemoryImage, MemoryImageSlot, ModuleMemoryImag
 /// lifetime of this store or the Send/Sync-ness of this store. All of that must
 /// be respected by embedders (e.g. the `wasmtime::Store` structure). The theory
 /// is that `wasmtime::Store` handles all this correctly.
-pub unsafe trait Store {
-    /// Returns the raw pointer in memory where this store's shared
-    /// `VMRuntimeLimits` structure is located.
-    ///
-    /// Used to configure `VMContext` initialization and store the right pointer
-    /// in the `VMContext`.
-    fn vmruntime_limits(&self) -> *mut VMRuntimeLimits;
-
-    /// Returns a pointer to the global epoch counter.
-    ///
-    /// Used to configure the `VMContext` on initialization.
-    fn epoch_ptr(&self) -> *const AtomicU64;
+pub unsafe trait VMStore {
+    /// Get a shared borrow of this store's `StoreOpaque`.
+    fn store_opaque(&self) -> &StoreOpaque;
 
-    /// Get this store's GC heap.
-    fn gc_store(&mut self) -> &mut GcStore {
-        self.maybe_gc_store()
-            .expect("attempt to access the GC store before it has been allocated")
-    }
-
-    /// Get this store's GC heap, if it has been allocated.
-    fn maybe_gc_store(&mut self) -> Option<&mut GcStore>;
+    /// Get an exclusive borrow of this store's `StoreOpaque`.
+    fn store_opaque_mut(&mut self) -> &mut StoreOpaque;
 
     /// Callback invoked to allow the store's resource limiter to reject a
     /// memory grow operation.
@@ -170,6 +158,20 @@ pub unsafe trait Store {
     fn component_calls(&mut self) -> &mut component::CallContexts;
 }
 
+impl Deref for dyn VMStore {
+    type Target = StoreOpaque;
+
+    fn deref(&self) -> &Self::Target {
+        self.store_opaque()
+    }
+}
+
+impl DerefMut for dyn VMStore {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.store_opaque_mut()
+    }
+}
+
 /// Functionality required by this crate for a particular module. This
 /// is chiefly needed for lazy initialization of various bits of
 /// instance state.

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

@@ -8,8 +8,8 @@
 
 use crate::prelude::*;
 use crate::runtime::vm::{
-    SendSyncPtr, Store, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
-    VMOpaqueContext, VMWasmCallFunction, ValRaw,
+    SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
+    VMOpaqueContext, VMStore, VMWasmCallFunction, ValRaw,
 };
 use alloc::alloc::Layout;
 use alloc::sync::Arc;
@@ -189,7 +189,7 @@ impl ComponentInstance {
         offsets: VMComponentOffsets<HostPtr>,
         runtime_info: Arc<dyn ComponentRuntimeInfo>,
         resource_types: Arc<dyn Any + Send + Sync>,
-        store: *mut dyn Store,
+        store: *mut dyn VMStore,
     ) {
         assert!(alloc_size >= Self::alloc_layout(&offsets).size());
 
@@ -253,9 +253,9 @@ impl ComponentInstance {
     }
 
     /// Returns the store that this component was created with.
-    pub fn store(&self) -> *mut dyn Store {
+    pub fn store(&self) -> *mut dyn VMStore {
         unsafe {
-            let ret = *self.vmctx_plus_offset::<*mut dyn Store>(self.offsets.store());
+            let ret = *self.vmctx_plus_offset::<*mut dyn VMStore>(self.offsets.store());
             assert!(!ret.is_null());
             ret
         }
@@ -435,7 +435,7 @@ impl ComponentInstance {
         }
     }
 
-    unsafe fn initialize_vmctx(&mut self, store: *mut dyn Store) {
+    unsafe fn initialize_vmctx(&mut self, store: *mut dyn VMStore) {
         *self.vmctx_plus_offset_mut(self.offsets.magic()) = VMCOMPONENT_MAGIC;
         *self.vmctx_plus_offset_mut(self.offsets.libcalls()) = &libcalls::VMComponentLibcalls::INIT;
         *self.vmctx_plus_offset_mut(self.offsets.store()) = store;
@@ -658,7 +658,7 @@ impl OwnedComponentInstance {
     pub fn new(
         runtime_info: Arc<dyn ComponentRuntimeInfo>,
         resource_types: Arc<dyn Any + Send + Sync>,
-        store: *mut dyn Store,
+        store: *mut dyn VMStore,
     ) -> OwnedComponentInstance {
         let component = runtime_info.component();
         let offsets = VMComponentOffsets::new(HostPtr, component);

crates/wasmtime/src/runtime/vm/const_expr.rs

@@ -28,7 +28,7 @@ impl<'a, 'b> ConstEvalContext<'a, 'b> {
 
     fn global_get(&mut self, index: GlobalIndex) -> Result<ValRaw> {
         unsafe {
-            let gc_store = (*self.instance.store()).gc_store();
+            let gc_store = (*self.instance.store()).unwrap_gc_store_mut();
             let global = self
                 .instance
                 .defined_or_imported_global_ptr(index)

crates/wasmtime/src/runtime/vm/instance.rs

@@ -14,7 +14,7 @@ use crate::runtime::vm::vmcontext::{
 };
 use crate::runtime::vm::{
     ExportFunction, ExportGlobal, ExportMemory, ExportTable, GcStore, Imports, ModuleRuntimeInfo,
-    SendSyncPtr, Store, VMFunctionBody, VMGcRef, WasmFault,
+    SendSyncPtr, VMFunctionBody, VMGcRef, VMStore, WasmFault,
 };
 use alloc::sync::Arc;
 use core::alloc::Layout;
@@ -461,22 +461,23 @@ impl Instance {
     /// functions are shared amongst threads and don't all share the same
     /// store).
     #[inline]
-    pub fn store(&self) -> *mut dyn Store {
-        let ptr =
-            unsafe { *self.vmctx_plus_offset::<*mut dyn Store>(self.offsets().ptr.vmctx_store()) };
+    pub fn store(&self) -> *mut dyn VMStore {
+        let ptr = unsafe {
+            *self.vmctx_plus_offset::<*mut dyn VMStore>(self.offsets().ptr.vmctx_store())
+        };
         debug_assert!(!ptr.is_null());
         ptr
     }
 
-    pub(crate) unsafe fn set_store(&mut self, store: Option<*mut dyn Store>) {
+    pub(crate) unsafe fn set_store(&mut self, store: Option<*mut dyn VMStore>) {
         if let Some(store) = store {
             *self.vmctx_plus_offset_mut(self.offsets().ptr.vmctx_store()) = store;
             *self.runtime_limits() = (*store).vmruntime_limits();
-            *self.epoch_ptr() = (*store).epoch_ptr();
-            self.set_gc_heap((*store).maybe_gc_store());
+            *self.epoch_ptr() = (*store).engine().epoch_counter();
+            self.set_gc_heap((*store).gc_store_mut().ok());
         } else {
             assert_eq!(
-                mem::size_of::<*mut dyn Store>(),
+                mem::size_of::<*mut dyn VMStore>(),
                 mem::size_of::<[*mut (); 2]>()
             );
             *self.vmctx_plus_offset_mut::<[*mut (); 2]>(self.offsets().ptr.vmctx_store()) =
@@ -1102,7 +1103,7 @@ impl Instance {
 
         if elt_ty == TableElementType::Func {
             for i in range {
-                let gc_store = unsafe { (*self.store()).gc_store() };
+                let gc_store = unsafe { (*self.store()).unwrap_gc_store_mut() };
                 let value = match self.tables[idx].1.get(gc_store, i) {
                     Some(value) => value,
                     None => {
@@ -1480,15 +1481,15 @@ impl InstanceHandle {
 
     /// Returns the `Store` pointer that was stored on creation
     #[inline]
-    pub fn store(&self) -> *mut dyn Store {
+    pub fn store(&self) -> *mut dyn VMStore {
         self.instance().store()
     }
 
     /// Configure the `*mut dyn Store` internal pointer after-the-fact.
     ///
     /// This is provided for the original `Store` itself to configure the first
     /// self-pointer after the original `Box` has been initialized.
-    pub unsafe fn set_store(&mut self, store: *mut dyn Store) {
+    pub unsafe fn set_store(&mut self, store: *mut dyn VMStore) {
         self.instance_mut().set_store(Some(store));
     }
 

crates/wasmtime/src/runtime/vm/instance/allocator.rs

@@ -5,7 +5,7 @@ use crate::runtime::vm::instance::{Instance, InstanceHandle};
 use crate::runtime::vm::memory::Memory;
 use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::table::Table;
-use crate::runtime::vm::{CompiledModuleId, ModuleRuntimeInfo, Store, VMFuncRef, VMGcRef};
+use crate::runtime::vm::{CompiledModuleId, ModuleRuntimeInfo, VMFuncRef, VMGcRef, VMStore};
 use core::{any::Any, mem, ptr};
 use wasmtime_environ::{
     DefinedMemoryIndex, DefinedTableIndex, HostPtr, InitMemory, MemoryInitialization,
@@ -79,7 +79,7 @@ pub struct InstanceAllocationRequest<'a> {
 /// InstanceAllocationRequest, rather than on a &mut InstanceAllocationRequest
 /// itself, because several use-sites require a split mut borrow on the
 /// InstanceAllocationRequest.
-pub struct StorePtr(Option<*mut dyn Store>);
+pub struct StorePtr(Option<*mut dyn VMStore>);
 
 impl StorePtr {
     /// A pointer to no Store.
@@ -88,19 +88,19 @@ impl StorePtr {
     }
 
     /// A pointer to a Store.
-    pub fn new(ptr: *mut dyn Store) -> Self {
+    pub fn new(ptr: *mut dyn VMStore) -> Self {
         Self(Some(ptr))
     }
 
     /// The raw contents of this struct
-    pub fn as_raw(&self) -> Option<*mut dyn Store> {
+    pub fn as_raw(&self) -> Option<*mut dyn VMStore> {
         self.0
     }
 
     /// Use the StorePtr as a mut ref to the Store.
     ///
     /// Safety: must not be used outside the original lifetime of the borrow.
-    pub(crate) unsafe fn get(&mut self) -> Option<&mut dyn Store> {
+    pub(crate) unsafe fn get(&mut self) -> Option<&mut dyn VMStore> {
         match self.0 {
             Some(ptr) => Some(&mut *ptr),
             None => None,
@@ -575,15 +575,15 @@ fn initialize_tables(instance: &mut Instance, module: &Module) -> Result<()> {
                 match module.table_plans[idx].table.ref_type.heap_type.top() {
                     WasmHeapTopType::Extern => {
                         let gc_ref = VMGcRef::from_raw_u32(raw.get_externref());
-                        let gc_store = unsafe { (*instance.store()).gc_store() };
+                        let gc_store = unsafe { (*instance.store()).unwrap_gc_store_mut() };
                         let items = (0..table.size())
                             .map(|_| gc_ref.as_ref().map(|r| gc_store.clone_gc_ref(r)));
                         table.init_gc_refs(0, items).err2anyhow()?;
                     }
 
                     WasmHeapTopType::Any => {
                         let gc_ref = VMGcRef::from_raw_u32(raw.get_anyref());
-                        let gc_store = unsafe { (*instance.store()).gc_store() };
+                        let gc_store = unsafe { (*instance.store()).unwrap_gc_store_mut() };
                         let items = (0..table.size())
                             .map(|_| gc_ref.as_ref().map(|r| gc_store.clone_gc_ref(r)));
                         table.init_gc_refs(0, items).err2anyhow()?;