Refactor component host/libcalls (#10959)

In the spirit of making Wasmtime's internals safer this is a step
forward for components to a new paradigm for how libcalls/host functions
are implemented. Previously `*mut ComponentInstance` was liberally used
but this meant that situations would often simultaneously have `&mut
ComponentInstance` and `&mut StoreOpaque` accessible in the same
function and there was no prevention of going from the store to the
component instance, acquiring two aliasing mutable references (which
would be unsound). The refactoring applied here is to redefine the
entrypoints from the guest back into the host to operate on `&mut dyn
VMStore` (or `StoreContextMut<'_, T>`) plus
`wasmtime::component::Instance`. This index-based approach means that
there's no aliasing of component instances and stores and the `Instance`
type can be used to look up anything within the store that's necessary.

This refactoring originated in the wasip3-prototyping repository and has
been used to remove a good deal of `unsafe` code now that `Instance` is
effectively safe to pass around and the store was already passed around
anyway everywhere.

In the future I plan to apply a similar paradigm shift for core
instances as well, but that'll require some more finesse for all the
bits and bobs that core wasm does.

Author: alexcrichton

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

@@ -4,7 +4,7 @@ use crate::component::storage::slice_to_storage_mut;
 use crate::component::{ComponentNamedList, ComponentType, Instance, Lift, Lower, Val};
 use crate::prelude::*;
 use crate::runtime::vm::component::{
-    InstanceFlags, VMComponentContext, VMLowering, VMLoweringCallee,
+    ComponentInstance, InstanceFlags, VMComponentContext, VMLowering, VMLoweringCallee,
 };
 use crate::runtime::vm::{VMFuncRef, VMGlobalDefinition, VMMemoryDefinition, VMOpaqueContext};
 use crate::{AsContextMut, CallHook, StoreContextMut, ValRaw};
@@ -13,8 +13,8 @@ use core::any::Any;
 use core::mem::{self, MaybeUninit};
 use core::ptr::NonNull;
 use wasmtime_environ::component::{
-    CanonicalAbiInfo, ComponentTypes, InterfaceType, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS,
-    StringEncoding, TypeFuncIndex,
+    CanonicalAbiInfo, InterfaceType, MAX_FLAT_PARAMS, MAX_FLAT_RESULTS, StringEncoding,
+    TypeFuncIndex,
 };
 
 pub struct HostFunc {
@@ -66,11 +66,10 @@ impl HostFunc {
     {
         let data = data.as_ptr() as *const F;
         unsafe {
-            call_host_and_handle_result::<T>(cx, |store, instance, types| {
+            call_host_and_handle_result::<T>(cx, |store, instance| {
                 call_host(
                     store,
                     instance,
-                    types,
                     TypeFuncIndex::from_u32(ty),
                     InstanceFlags::from_raw(flags),
                     memory,
@@ -148,7 +147,6 @@ where
 unsafe fn call_host<T, Params, Return, F>(
     mut cx: StoreContextMut<'_, T>,
     instance: Instance,
-    types: &Arc<ComponentTypes>,
     ty: TypeFuncIndex,
     mut flags: InstanceFlags,
     memory: *mut VMMemoryDefinition,
@@ -199,6 +197,7 @@ where
         bail!("cannot leave component instance");
     }
 
+    let types = cx[instance.id()].component().types().clone();
     let ty = &types[ty];
     let param_tys = InterfaceType::Tuple(ty.params);
     let result_tys = InterfaceType::Tuple(ty.results);
@@ -225,13 +224,13 @@ where
             Storage::Indirect(slice_to_storage_mut(storage).assume_init_ref())
         }
     };
-    let mut lift = LiftContext::new(cx.0, &options, types, instance);
+    let mut lift = LiftContext::new(cx.0, &options, &types, instance);
     lift.enter_call();
     let params = storage.lift_params(&mut lift, param_tys)?;
 
     let ret = closure(cx.as_context_mut(), params)?;
     flags.set_may_leave(false);
-    let mut lower = LowerContext::new(cx, &options, types, instance);
+    let mut lower = LowerContext::new(cx, &options, &types, instance);
     storage.lower_results(&mut lower, result_tys, ret)?;
     flags.set_may_leave(true);
 
@@ -308,30 +307,27 @@ fn validate_inbounds<T: ComponentType>(memory: &[u8], ptr: &ValRaw) -> Result<us
 
 unsafe fn call_host_and_handle_result<T>(
     cx: NonNull<VMOpaqueContext>,
-    func: impl FnOnce(StoreContextMut<'_, T>, Instance, &Arc<ComponentTypes>) -> Result<()>,
+    func: impl FnOnce(StoreContextMut<'_, T>, Instance) -> Result<()>,
 ) -> bool
 where
     T: 'static,
 {
     let cx = VMComponentContext::from_opaque(cx);
-    let instance = cx.as_ref().instance();
-    let types = (*instance).component().types();
-    let raw_store = (*instance).store();
-    let mut store = StoreContextMut(&mut *raw_store.cast());
-    let instance = Instance::from_wasmtime(store.0, (*instance).id());
-
-    crate::runtime::vm::catch_unwind_and_record_trap(|| {
-        store.0.call_hook(CallHook::CallingHost)?;
-        let res = func(store.as_context_mut(), instance, types);
-        store.0.call_hook(CallHook::ReturningFromHost)?;
-        res
+    ComponentInstance::from_vmctx(cx, |store, instance| {
+        let mut store = store.unchecked_context_mut();
+
+        crate::runtime::vm::catch_unwind_and_record_trap(|| {
+            store.0.call_hook(CallHook::CallingHost)?;
+            let res = func(store.as_context_mut(), instance);
+            store.0.call_hook(CallHook::ReturningFromHost)?;
+            res
+        })
     })
 }
 
 unsafe fn call_host_dynamic<T, F>(
     mut store: StoreContextMut<'_, T>,
     instance: Instance,
-    types: &Arc<ComponentTypes>,
     ty: TypeFuncIndex,
     mut flags: InstanceFlags,
     memory: *mut VMMemoryDefinition,
@@ -366,10 +362,11 @@ where
     let args;
     let ret_index;
 
+    let types = store[instance.id()].component().types().clone();
     let func_ty = &types[ty];
     let param_tys = &types[func_ty.params];
     let result_tys = &types[func_ty.results];
-    let mut cx = LiftContext::new(store.0, &options, types, instance);
+    let mut cx = LiftContext::new(store.0, &options, &types, instance);
     cx.enter_call();
     if let Some(param_count) = param_tys.abi.flat_count(MAX_FLAT_PARAMS) {
         // NB: can use `MaybeUninit::slice_assume_init_ref` when that's stable
@@ -405,7 +402,7 @@ where
     closure(store.as_context_mut(), &args, &mut result_vals)?;
     flags.set_may_leave(false);
 
-    let mut cx = LowerContext::new(store, &options, types, instance);
+    let mut cx = LowerContext::new(store, &options, &types, instance);
     if let Some(cnt) = result_tys.abi.flat_count(MAX_FLAT_RESULTS) {
         let mut dst = storage[..cnt].iter_mut();
         for (val, ty) in result_vals.iter().zip(result_tys.types.iter()) {
@@ -463,11 +460,10 @@ where
 {
     let data = data.as_ptr() as *const F;
     unsafe {
-        call_host_and_handle_result(cx, |store, instance, types| {
+        call_host_and_handle_result(cx, |store, instance| {
             call_host_dynamic::<T, _>(
                 store,
                 instance,
-                types,
                 TypeFuncIndex::from_u32(ty),
                 InstanceFlags::from_raw(flags),
                 memory,

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

@@ -9,7 +9,10 @@ use crate::instance::OwnedImports;
 use crate::linker::DefinitionType;
 use crate::prelude::*;
 use crate::runtime::vm::VMFuncRef;
-use crate::runtime::vm::component::{ComponentInstance, OwnedComponentInstance};
+use crate::runtime::vm::component::{
+    CallContexts, ComponentInstance, OwnedComponentInstance, ResourceTables, TypedResource,
+    TypedResourceIndex,
+};
 use crate::store::StoreOpaque;
 use crate::{AsContext, AsContextMut, Engine, Module, StoreContextMut};
 use alloc::sync::Arc;
@@ -376,6 +379,100 @@ impl Instance {
     pub(crate) fn id(&self) -> StoreComponentInstanceId {
         self.id
     }
+
+    /// Implementation of the `resource.new` intrinsic for `i32`
+    /// representations.
+    pub(crate) fn resource_new32(
+        self,
+        store: &mut StoreOpaque,
+        ty: TypeResourceTableIndex,
+        rep: u32,
+    ) -> Result<u32> {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        resource_tables(calls, instance).resource_new(TypedResource::Component { ty, rep })
+    }
+
+    /// Implementation of the `resource.rep` intrinsic for `i32`
+    /// representations.
+    pub(crate) fn resource_rep32(
+        self,
+        store: &mut StoreOpaque,
+        ty: TypeResourceTableIndex,
+        index: u32,
+    ) -> Result<u32> {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        resource_tables(calls, instance).resource_rep(TypedResourceIndex::Component { ty, index })
+    }
+
+    /// Implementation of the `resource.drop` intrinsic.
+    pub(crate) fn resource_drop(
+        self,
+        store: &mut StoreOpaque,
+        ty: TypeResourceTableIndex,
+        index: u32,
+    ) -> Result<Option<u32>> {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        resource_tables(calls, instance).resource_drop(TypedResourceIndex::Component { ty, index })
+    }
+
+    pub(crate) fn resource_transfer_own(
+        self,
+        store: &mut StoreOpaque,
+        index: u32,
+        src: TypeResourceTableIndex,
+        dst: TypeResourceTableIndex,
+    ) -> Result<u32> {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        let mut tables = resource_tables(calls, instance);
+        let rep = tables.resource_lift_own(TypedResourceIndex::Component { ty: src, index })?;
+        tables.resource_lower_own(TypedResource::Component { ty: dst, rep })
+    }
+
+    pub(crate) fn resource_transfer_borrow(
+        self,
+        store: &mut StoreOpaque,
+        index: u32,
+        src: TypeResourceTableIndex,
+        dst: TypeResourceTableIndex,
+    ) -> Result<u32> {
+        let dst_owns_resource = store[self.id()].resource_owned_by_own_instance(dst);
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        let mut tables = resource_tables(calls, instance);
+        let rep = tables.resource_lift_borrow(TypedResourceIndex::Component { ty: src, index })?;
+        // Implement `lower_borrow`'s special case here where if a borrow's
+        // resource type is owned by `dst` then the destination receives the
+        // representation directly rather than a handle to the representation.
+        //
+        // This can perhaps become a different libcall in the future to avoid
+        // this check at runtime since we know at compile time whether the
+        // destination type owns the resource, but that's left as a future
+        // refactoring if truly necessary.
+        if dst_owns_resource {
+            return Ok(rep);
+        }
+        tables.resource_lower_borrow(TypedResource::Component { ty: dst, rep })
+    }
+
+    pub(crate) fn resource_enter_call(self, store: &mut StoreOpaque) {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        resource_tables(calls, instance).enter_call()
+    }
+
+    pub(crate) fn resource_exit_call(self, store: &mut StoreOpaque) -> Result<()> {
+        let (calls, _, _, instance) = store.component_resource_state_with_instance(self);
+        resource_tables(calls, instance).exit_call()
+    }
+}
+
+fn resource_tables<'a>(
+    calls: &'a mut CallContexts,
+    instance: &'a mut ComponentInstance,
+) -> ResourceTables<'a> {
+    ResourceTables {
+        host_table: None,
+        calls,
+        guest: Some(instance.guest_tables()),
+    }
 }
 
 /// Trait used to lookup the export of a component instance.

crates/wasmtime/src/runtime/externals/global.rs

@@ -287,7 +287,9 @@ impl Global {
             ),
             #[cfg(feature = "component-model")]
             ExportGlobalKind::ComponentFlags(vmctx, index) => (
-                vm::component::ComponentInstance::from_vmctx(vmctx, |i| i.id().as_u32()),
+                vm::component::ComponentInstance::from_vmctx(vmctx, |_, i| {
+                    i.id().instance().as_u32()
+                }),
                 VMGlobalKind::ComponentFlags(index),
             ),
         };

crates/wasmtime/src/runtime/func.rs

@@ -2048,7 +2048,8 @@ impl<T> Caller<'_, T> {
         R: 'static,
     {
         crate::runtime::vm::InstanceAndStore::from_vmctx(caller, |pair| {
-            let (instance, mut store) = pair.unpack_context_mut::<T>();
+            let (instance, store) = pair.unpack_mut();
+            let mut store = store.unchecked_context_mut::<T>();
             let caller = Instance::from_wasmtime(instance.id(), store.0);
 
             let (gc_lifo_scope, ret) = {

crates/wasmtime/src/runtime/vm.rs

@@ -33,7 +33,9 @@ pub(crate) struct f32x4(crate::uninhabited::Uninhabited);
 #[derive(Copy, Clone)]
 pub(crate) struct f64x2(crate::uninhabited::Uninhabited);
 
+use crate::StoreContextMut;
 use crate::prelude::*;
+use crate::store::StoreInner;
 use crate::store::StoreOpaque;
 use alloc::sync::Arc;
 use core::fmt;
@@ -166,13 +168,21 @@ cfg_if::cfg_if! {
 ///
 /// This trait is used to store a raw pointer trait object within each
 /// `VMContext`. This raw pointer trait object points back to the
-/// `wasmtime::Store` internally but is type-erased so this `wasmtime-runtime`
-/// crate doesn't need the entire `wasmtime` crate to build.
+/// `wasmtime::Store` internally but is type-erased to avoid needing to
+/// monomorphize the entire runtime on the `T` in `Store<T>`
 ///
-/// Note that this is an extra-unsafe trait because no heed is paid to the
-/// 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.
+/// # Safety
+///
+/// This trait should be implemented by nothing other than `StoreInner<T>` in
+/// this crate. It's not sound to implement it for anything else due to
+/// `unchecked_context_mut` below.
+///
+/// It's also worth nothing that there are various locations where a `*mut dyn
+/// VMStore` is asserted to be both `Send` and `Sync` which disregards the `T`
+/// that's actually stored in the store itself. It's assume that the high-level
+/// APIs using `Store<T>` are correctly inferring send/sync on the returned
+/// values (e.g. futures) and that internally in the runtime we aren't doing
+/// anything "weird" with threads for example.
 pub unsafe trait VMStore {
     /// Get a shared borrow of this store's `StoreOpaque`.
     fn store_opaque(&self) -> &StoreOpaque;
@@ -265,6 +275,19 @@ impl DerefMut for dyn VMStore + '_ {
     }
 }
 
+impl dyn VMStore + '_ {
+    /// Asserts that this `VMStore` was originally paired with `StoreInner<T>`
+    /// and then casts to the `StoreContextMut` type.
+    ///
+    /// # Unsafety
+    ///
+    /// This method is not safe as there's no static guarantee that `T` is
+    /// correct for this store.
+    pub(crate) unsafe fn unchecked_context_mut<T>(&mut self) -> StoreContextMut<'_, T> {
+        StoreContextMut(&mut *(self as *mut dyn VMStore as *mut StoreInner<T>))
+    }
+}
+
 /// A newtype wrapper around `NonNull<dyn VMStore>` intended to be a
 /// self-pointer back to the `Store<T>` within raw data structures like
 /// `VMContext`.