Move *mut dyn VMStore out of VMContext (#10040)

* Move `*mut dyn VMStore` out of `VMContext`

There's no need for this type to be accessible from compiled code, so
store and manage it entirely from Rust instead.

* Fix test expectations

* Update disas test expectations

Author: alexcrichton

crates/environ/src/component/vmcomponent_offsets.rs

@@ -3,7 +3,6 @@
 // struct VMComponentContext {
 //      magic: u32,
 //      builtins: &'static VMComponentBuiltins,
-//      store: *mut dyn Store,
 //      limits: *const VMRuntimeLimits,
 //      flags: [VMGlobalDefinition; component.num_runtime_component_instances],
 //      trampoline_func_refs: [VMFuncRef; component.num_trampolines],
@@ -62,7 +61,6 @@ pub struct VMComponentOffsets<P> {
     // precalculated offsets of various member fields
     magic: u32,
     builtins: u32,
-    store: u32,
     limits: u32,
     flags: u32,
     trampoline_func_refs: u32,
@@ -100,7 +98,6 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             num_resources: component.num_resources,
             magic: 0,
             builtins: 0,
-            store: 0,
             limits: 0,
             flags: 0,
             trampoline_func_refs: 0,
@@ -141,7 +138,6 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             size(magic) = 4u32,
             align(u32::from(ret.ptr.size())),
             size(builtins) = ret.ptr.size(),
-            size(store) = cmul(2, ret.ptr.size()),
             size(limits) = ret.ptr.size(),
             align(16),
             size(flags) = cmul(ret.num_runtime_component_instances, ret.ptr.size_of_vmglobal_definition()),
@@ -190,12 +186,6 @@ impl<P: PtrSize> VMComponentOffsets<P> {
         self.flags + index.as_u32() * u32::from(self.ptr.size_of_vmglobal_definition())
     }
 
-    /// The offset of the `store` field.
-    #[inline]
-    pub fn store(&self) -> u32 {
-        self.store
-    }
-
     /// The offset of the `limits` field.
     #[inline]
     pub fn limits(&self) -> u32 {

crates/environ/src/vmoffsets.rs

@@ -15,7 +15,6 @@
 //      gc_heap_base: *mut u8,
 //      gc_heap_bound: *mut u8,
 //      gc_heap_data: *mut T, // Collector-specific pointer
-//      store: *mut dyn Store,
 //      type_ids: *const VMSharedTypeIndex,
 //
 //      // Variable-width fields come after the fixed-width fields above. Place
@@ -278,16 +277,10 @@ pub trait PtrSize {
         self.vmctx_gc_heap_bound() + self.size()
     }
 
-    /// The offset of the `*const dyn Store` member.
-    #[inline]
-    fn vmctx_store(&self) -> u8 {
-        self.vmctx_gc_heap_data() + self.size()
-    }
-
     /// The offset of the `type_ids` array pointer.
     #[inline]
     fn vmctx_type_ids_array(&self) -> u8 {
-        self.vmctx_store() + 2 * self.size()
+        self.vmctx_gc_heap_data() + self.size()
     }
 
     /// The end of statically known offsets in `VMContext`.

crates/wasmtime/src/runtime/store.rs

@@ -103,6 +103,7 @@ use core::num::NonZeroU64;
 use core::ops::{Deref, DerefMut, Range};
 use core::pin::Pin;
 use core::ptr;
+use core::ptr::NonNull;
 use core::task::{Context, Poll};
 use wasmtime_environ::TripleExt;
 
@@ -628,9 +629,9 @@ impl<T> Store<T> {
             // maintain throughout Wasmtime.
             unsafe {
                 let traitobj = mem::transmute::<
-                    *mut (dyn crate::runtime::vm::VMStore + '_),
-                    *mut (dyn crate::runtime::vm::VMStore + 'static),
-                >(&mut *inner);
+                    NonNull<dyn crate::runtime::vm::VMStore + '_>,
+                    NonNull<dyn crate::runtime::vm::VMStore + 'static>,
+                >(NonNull::from(&mut *inner));
                 instance.set_store(traitobj);
                 instance
             }
@@ -1933,7 +1934,7 @@ impl StoreOpaque {
     }
 
     #[inline]
-    pub fn traitobj(&self) -> *mut dyn crate::runtime::vm::VMStore {
+    pub fn traitobj(&self) -> NonNull<dyn crate::runtime::vm::VMStore> {
         self.default_caller.traitobj(self)
     }
 

crates/wasmtime/src/runtime/vm.rs

@@ -203,6 +203,29 @@ impl DerefMut for dyn VMStore + '_ {
     }
 }
 
+/// A newtype wrapper around `NonNull<dyn VMStore>` intended to be a
+/// self-pointer back to the `Store<T>` within raw data structures like
+/// `VMContext`.
+///
+/// This type exists to manually, and unsafely, implement `Send` and `Sync`.
+/// The `VMStore` trait doesn't require `Send` or `Sync` which means this isn't
+/// naturally either trait (e.g. with `SendSyncPtr` instead). Note that this
+/// means that `Instance` is, for example, mistakenly considered
+/// unconditionally `Send` and `Sync`. This is hopefully ok for now though
+/// because from a user perspective the only type that matters is `Store<T>`.
+/// That type is `Send + Sync` if `T: Send + Sync` already so the internal
+/// storage of `Instance` shouldn't matter as the final result is the same.
+/// Note though that this means we need to be extra vigilant about cross-thread
+/// usage of `Instance` and `ComponentInstance` for example.
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+struct VMStoreRawPtr(NonNull<dyn VMStore>);
+
+// SAFETY: this is the purpose of `VMStoreRawPtr`, see docs above about safe
+// usage.
+unsafe impl Send for VMStoreRawPtr {}
+unsafe impl Sync for VMStoreRawPtr {}
+
 /// 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

@@ -9,7 +9,7 @@
 use crate::prelude::*;
 use crate::runtime::vm::{
     SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
-    VMOpaqueContext, VMStore, VMWasmCallFunction, ValRaw,
+    VMOpaqueContext, VMStore, VMStoreRawPtr, VMWasmCallFunction, ValRaw,
 };
 use alloc::alloc::Layout;
 use alloc::sync::Arc;
@@ -66,6 +66,9 @@ pub struct ComponentInstance {
     /// Any` is left as an exercise for a future refactoring.
     resource_types: Arc<dyn Any + Send + Sync>,
 
+    /// Self-pointer back to `Store<T>` and its functions.
+    store: VMStoreRawPtr,
+
     /// A zero-sized field which represents the end of the struct for the actual
     /// `VMComponentContext` to be allocated behind.
     vmctx: VMComponentContext,
@@ -193,7 +196,7 @@ impl ComponentInstance {
         offsets: VMComponentOffsets<HostPtr>,
         runtime_info: Arc<dyn ComponentRuntimeInfo>,
         resource_types: Arc<dyn Any + Send + Sync>,
-        store: *mut dyn VMStore,
+        store: NonNull<dyn VMStore>,
     ) {
         assert!(alloc_size >= Self::alloc_layout(&offsets).size());
 
@@ -218,13 +221,14 @@ impl ComponentInstance {
                 component_resource_tables,
                 runtime_info,
                 resource_types,
+                store: VMStoreRawPtr(store),
                 vmctx: VMComponentContext {
                     _marker: marker::PhantomPinned,
                 },
             },
         );
 
-        (*ptr.as_ptr()).initialize_vmctx(store);
+        (*ptr.as_ptr()).initialize_vmctx();
     }
 
     fn vmctx(&self) -> *mut VMComponentContext {
@@ -258,11 +262,7 @@ impl ComponentInstance {
 
     /// Returns the store that this component was created with.
     pub fn store(&self) -> *mut dyn VMStore {
-        unsafe {
-            let ret = *self.vmctx_plus_offset::<*mut dyn VMStore>(self.offsets.store());
-            assert!(!ret.is_null());
-            ret
-        }
+        self.store.0.as_ptr()
     }
 
     /// Returns the runtime memory definition corresponding to the index of the
@@ -439,11 +439,11 @@ impl ComponentInstance {
         }
     }
 
-    unsafe fn initialize_vmctx(&mut self, store: *mut dyn VMStore) {
+    unsafe fn initialize_vmctx(&mut self) {
         *self.vmctx_plus_offset_mut(self.offsets.magic()) = VMCOMPONENT_MAGIC;
         *self.vmctx_plus_offset_mut(self.offsets.builtins()) = &libcalls::VMComponentBuiltins::INIT;
-        *self.vmctx_plus_offset_mut(self.offsets.store()) = store;
-        *self.vmctx_plus_offset_mut(self.offsets.limits()) = (*store).vmruntime_limits();
+        *self.vmctx_plus_offset_mut(self.offsets.limits()) =
+            self.store.0.as_ref().vmruntime_limits();
 
         for i in 0..self.offsets.num_runtime_component_instances {
             let i = RuntimeComponentInstanceIndex::from_u32(i);
@@ -662,7 +662,7 @@ impl OwnedComponentInstance {
     pub fn new(
         runtime_info: Arc<dyn ComponentRuntimeInfo>,
         resource_types: Arc<dyn Any + Send + Sync>,
-        store: *mut dyn VMStore,
+        store: NonNull<dyn VMStore>,
     ) -> OwnedComponentInstance {
         let component = runtime_info.component();
         let offsets = VMComponentOffsets::new(HostPtr, component);

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

@@ -13,7 +13,7 @@ use crate::runtime::vm::vmcontext::{
 };
 use crate::runtime::vm::{
     ExportFunction, ExportGlobal, ExportMemory, ExportTable, GcStore, Imports, ModuleRuntimeInfo,
-    SendSyncPtr, VMFunctionBody, VMGcRef, VMStore, WasmFault,
+    SendSyncPtr, VMFunctionBody, VMGcRef, VMStore, VMStoreRawPtr, WasmFault,
 };
 use crate::store::{StoreInner, StoreOpaque};
 use crate::{prelude::*, StoreContextMut};
@@ -162,13 +162,7 @@ impl InstanceAndStore {
     /// store).
     #[inline]
     fn store_ptr(&self) -> *mut dyn VMStore {
-        let ptr = unsafe {
-            *self
-                .instance
-                .vmctx_plus_offset::<*mut dyn VMStore>(self.instance.offsets().ptr.vmctx_store())
-        };
-        debug_assert!(!ptr.is_null());
-        ptr
+        self.instance.store.unwrap().0.as_ptr()
     }
 }
 
@@ -281,6 +275,12 @@ pub struct Instance {
     #[cfg(feature = "wmemcheck")]
     pub(crate) wmemcheck_state: Option<Wmemcheck>,
 
+    /// Self-pointer back to `Store<T>` and its functions. Not present for
+    /// the brief time that `Store<T>` is itself being created. Also not
+    /// present for some niche uses that are disconnected from stores (e.g.
+    /// cross-thread stuff used in `InstancePre`)
+    store: Option<VMStoreRawPtr>,
+
     /// Additional context used by compiled wasm code. This field is last, and
     /// represents a dynamically-sized array that extends beyond the nominal
     /// end of the struct (similar to a flexible array member).
@@ -341,6 +341,7 @@ impl Instance {
                         None
                     }
                 },
+                store: None,
             },
         );
 
@@ -584,19 +585,14 @@ impl Instance {
         unsafe { self.vmctx_plus_offset_mut(self.offsets().ptr.vmctx_gc_heap_data()) }
     }
 
-    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).engine().epoch_counter();
-            self.set_gc_heap((*store).gc_store_mut().ok());
+    pub(crate) unsafe fn set_store(&mut self, store: Option<NonNull<dyn VMStore>>) {
+        self.store = store.map(VMStoreRawPtr);
+        if let Some(mut store) = store {
+            let store = store.as_mut();
+            *self.runtime_limits() = store.vmruntime_limits();
+            *self.epoch_ptr() = store.engine().epoch_counter();
+            self.set_gc_heap(store.gc_store_mut().ok());
         } else {
-            assert_eq!(
-                mem::size_of::<*mut dyn VMStore>(),
-                mem::size_of::<[*mut (); 2]>()
-            );
-            *self.vmctx_plus_offset_mut::<[*mut (); 2]>(self.offsets().ptr.vmctx_store()) =
-                [ptr::null_mut(), ptr::null_mut()];
             *self.runtime_limits() = ptr::null_mut();
             *self.epoch_ptr() = ptr::null_mut();
             self.set_gc_heap(None);
@@ -1607,28 +1603,22 @@ impl InstanceHandle {
     /// This should only be used for initializing a vmctx's store pointer. It
     /// should never be used to access the store itself. Use `InstanceAndStore`
     /// for that instead.
-    pub fn traitobj(&self, store: &StoreOpaque) -> *mut dyn VMStore {
+    pub fn traitobj(&self, store: &StoreOpaque) -> NonNull<dyn VMStore> {
         // By requiring a store argument, we are ensuring that callers aren't
         // getting this trait object in order to access the store, since they
         // already have access. See `InstanceAndStore` and its documentation for
         // details about the store access patterns we want to restrict host code
         // to.
         let _ = store;
 
-        let ptr = unsafe {
-            *self
-                .instance()
-                .vmctx_plus_offset::<*mut dyn VMStore>(self.instance().offsets().ptr.vmctx_store())
-        };
-        debug_assert!(!ptr.is_null());
-        ptr
+        self.instance().store.unwrap().0
     }
 
     /// 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 VMStore) {
+    pub unsafe fn set_store(&mut self, store: NonNull<dyn VMStore>) {
         self.instance_mut().set_store(Some(store));
     }
 

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

@@ -85,7 +85,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 VMStore>);
+pub struct StorePtr(Option<NonNull<dyn VMStore>>);
 
 impl StorePtr {
     /// A pointer to no Store.
@@ -94,23 +94,21 @@ impl StorePtr {
     }
 
     /// A pointer to a Store.
-    pub fn new(ptr: *mut dyn VMStore) -> Self {
+    pub fn new(ptr: NonNull<dyn VMStore>) -> Self {
         Self(Some(ptr))
     }
 
     /// The raw contents of this struct
-    pub fn as_raw(&self) -> Option<*mut dyn VMStore> {
+    pub fn as_raw(&self) -> Option<NonNull<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 VMStore> {
-        match self.0 {
-            Some(ptr) => Some(&mut *ptr),
-            None => None,
-        }
+        let ptr = self.0?.as_mut();
+        Some(ptr)
     }
 }