feat: implement Resource<T> -> ResourceAny conversion (#7688)

* feat: implement `Resource<T>` -> `ResourceAny` conversion Signed-off-by: Roman Volosatovs <[email protected]> * feat(component/linker): introduce `PathIndex` Signed-off-by: Roman Volosatovs <[email protected]> * test: add `into_resource_any` tests Signed-off-by: Roman Volosatovs <[email protected]> * refactor(linker): only track import index for resources Signed-off-by: Roman Volosatovs <[email protected]> * feat(PrimaryMap): add `From<Vec<V>>` implementation Signed-off-by: Roman Volosatovs <[email protected]> * refactor: store `ResourceImportIndex` mapping in `PrimaryMap` Signed-off-by: Roman Volosatovs <[email protected]> * doc: expand on `ResourceImportIndex` purpose prtest:full Signed-off-by: Roman Volosatovs <[email protected]> * refactor: use `Linker::insert` in `Linker::resource` Signed-off-by: Roman Volosatovs <[email protected]> --------- Signed-off-by: Roman Volosatovs <[email protected]>
bytecodealliance · Jan 3, 2024 · ae5d284 · ae5d284
1 parent 2b2e638
commit ae5d284
Show file tree

Hide file tree

Showing 7 changed files with 218 additions and 20 deletions.
diff --git a/cranelift/entity/src/primary.rs b/cranelift/entity/src/primary.rs
@@ -306,6 +306,18 @@ where
     }
 }
 
+impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
+where
+    K: EntityRef,
+{
+    fn from(elems: Vec<V>) -> Self {
+        Self {
+            elems,
+            unused: PhantomData,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -496,6 +508,19 @@ mod tests {
         }
     }
 
+    #[test]
+    fn from_vec() {
+        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
+        m.push(12);
+        m.push(33);
+
+        let n = PrimaryMap::<E, &usize>::from(m.values().collect::<Vec<_>>());
+        assert!(m.len() == n.len());
+        for (me, ne) in m.values().zip(n.values()) {
+            assert!(*me == **ne);
+        }
+    }
+
     #[test]
     fn get_many_mut() {
         let mut m: PrimaryMap<E, usize> = PrimaryMap::new();

diff --git a/crates/wasmtime/src/component/instance.rs b/crates/wasmtime/src/component/instance.rs
@@ -1,6 +1,8 @@
 use crate::component::func::HostFunc;
 use crate::component::matching::InstanceType;
-use crate::component::{Component, ComponentNamedList, Func, Lift, Lower, ResourceType, TypedFunc};
+use crate::component::{
+    Component, ComponentNamedList, Func, Lift, Lower, ResourceImportIndex, ResourceType, TypedFunc,
+};
 use crate::instance::OwnedImports;
 use crate::linker::DefinitionType;
 use crate::store::{StoreOpaque, Stored};
@@ -520,6 +522,7 @@ impl<'a> Instantiator<'a> {
 pub struct InstancePre<T> {
     component: Component,
     imports: Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
+    resource_imports: Arc<PrimaryMap<ResourceImportIndex, Option<RuntimeImportIndex>>>,
     _marker: marker::PhantomData<fn() -> T>,
 }
 
@@ -529,6 +532,7 @@ impl<T> Clone for InstancePre<T> {
         Self {
             component: self.component.clone(),
             imports: self.imports.clone(),
+            resource_imports: self.resource_imports.clone(),
             _marker: self._marker,
         }
     }
@@ -544,14 +548,28 @@ impl<T> InstancePre<T> {
     pub(crate) unsafe fn new_unchecked(
         component: Component,
         imports: PrimaryMap<RuntimeImportIndex, RuntimeImport>,
+        resource_imports: PrimaryMap<ResourceImportIndex, Option<RuntimeImportIndex>>,
     ) -> InstancePre<T> {
         InstancePre {
             component,
             imports: Arc::new(imports),
+            resource_imports: Arc::new(resource_imports),
             _marker: marker::PhantomData,
         }
     }
 
+    pub(crate) fn resource_import_index(
+        &self,
+        path: ResourceImportIndex,
+    ) -> Option<RuntimeImportIndex> {
+        *self.resource_imports.get(path)?
+    }
+
+    pub(crate) fn resource_import(&self, path: ResourceImportIndex) -> Option<&RuntimeImport> {
+        let idx = self.resource_import_index(path)?;
+        self.imports.get(idx)
+    }
+
     /// Returns the underlying component that will be instantiated.
     pub fn component(&self) -> &Component {
         &self.component

diff --git a/crates/wasmtime/src/component/linker.rs b/crates/wasmtime/src/component/linker.rs
@@ -14,7 +14,7 @@ use std::ops::Deref;
 use std::pin::Pin;
 use std::sync::Arc;
 use wasmtime_environ::component::TypeDef;
-use wasmtime_environ::PrimaryMap;
+use wasmtime_environ::{EntityRef, PrimaryMap};
 
 /// A type used to instantiate [`Component`]s.
 ///
@@ -27,6 +27,7 @@ pub struct Linker<T> {
     strings: Strings,
     map: NameMap,
     path: Vec<usize>,
+    resource_imports: usize,
     allow_shadowing: bool,
     _marker: marker::PhantomData<fn() -> T>,
 }
@@ -38,6 +39,7 @@ impl<T> Clone for Linker<T> {
             strings: self.strings.clone(),
             map: self.map.clone(),
             path: self.path.clone(),
+            resource_imports: self.resource_imports.clone(),
             allow_shadowing: self.allow_shadowing,
             _marker: self._marker,
         }
@@ -61,18 +63,62 @@ pub struct LinkerInstance<'a, T> {
     path_len: usize,
     strings: &'a mut Strings,
     map: &'a mut NameMap,
+    resource_imports: &'a mut usize,
     allow_shadowing: bool,
     _marker: marker::PhantomData<fn() -> T>,
 }
 
+/// Index correlating a resource definition to the import path.
+/// This is assigned by [`Linker::resource`] and may be used to associate it to
+/// [`RuntimeImportIndex`](wasmtime_environ::component::RuntimeImportIndex)
+/// at a later stage
+///
+/// [`Linker::resource`]: crate::component::LinkerInstance::resource
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+pub struct ResourceImportIndex(usize);
+
+impl EntityRef for ResourceImportIndex {
+    fn new(idx: usize) -> Self {
+        Self(idx)
+    }
+
+    fn index(self) -> usize {
+        self.0
+    }
+}
+
+impl Deref for ResourceImportIndex {
+    type Target = usize;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl From<usize> for ResourceImportIndex {
+    fn from(idx: usize) -> Self {
+        Self(idx)
+    }
+}
+
+impl From<ResourceImportIndex> for usize {
+    fn from(idx: ResourceImportIndex) -> Self {
+        idx.0
+    }
+}
+
 pub(crate) type NameMap = HashMap<usize, Definition>;
 
 #[derive(Clone)]
 pub(crate) enum Definition {
     Instance(NameMap),
     Func(Arc<HostFunc>),
     Module(Module),
-    Resource(ResourceType, Arc<crate::func::HostFunc>),
+    Resource(
+        ResourceImportIndex,
+        ResourceType,
+        Arc<crate::func::HostFunc>,
+    ),
 }
 
 impl<T> Linker<T> {
@@ -85,6 +131,7 @@ impl<T> Linker<T> {
             map: NameMap::default(),
             allow_shadowing: false,
             path: Vec::new(),
+            resource_imports: 0,
             _marker: marker::PhantomData,
         }
     }
@@ -112,6 +159,7 @@ impl<T> Linker<T> {
             path_len: 0,
             strings: &mut self.strings,
             map: &mut self.map,
+            resource_imports: &mut self.resource_imports,
             allow_shadowing: self.allow_shadowing,
             _marker: self._marker,
         }
@@ -170,6 +218,7 @@ impl<T> Linker<T> {
         // using the import map within the component created at
         // component-compile-time.
         let mut imports = PrimaryMap::with_capacity(env_component.imports.len());
+        let mut resource_imports = PrimaryMap::from(vec![None; self.resource_imports]);
         for (idx, (import, names)) in env_component.imports.iter() {
             let (root, _) = &env_component.import_types[*import];
             let root = self.strings.lookup(root).unwrap();
@@ -188,11 +237,14 @@ impl<T> Linker<T> {
             let import = match cur {
                 Definition::Module(m) => RuntimeImport::Module(m.clone()),
                 Definition::Func(f) => RuntimeImport::Func(f.clone()),
-                Definition::Resource(t, dtor) => RuntimeImport::Resource {
-                    ty: t.clone(),
-                    _dtor: dtor.clone(),
-                    dtor_funcref: component.resource_drop_func_ref(dtor),
-                },
+                Definition::Resource(res_idx, t, dtor) => {
+                    resource_imports[*res_idx] = Some(idx);
+                    RuntimeImport::Resource {
+                        ty: t.clone(),
+                        _dtor: dtor.clone(),
+                        dtor_funcref: component.resource_drop_func_ref(dtor),
+                    }
+                }
 
                 // This is guaranteed by the compilation process that "leaf"
                 // runtime imports are never instances.
@@ -201,7 +253,7 @@ impl<T> Linker<T> {
             let i = imports.push(import);
             assert_eq!(i, idx);
         }
-        Ok(unsafe { InstancePre::new_unchecked(component.clone(), imports) })
+        Ok(unsafe { InstancePre::new_unchecked(component.clone(), imports, resource_imports) })
     }
 
     /// Instantiates the [`Component`] provided into the `store` specified.
@@ -266,6 +318,7 @@ impl<T> LinkerInstance<'_, T> {
             path_len: self.path_len,
             strings: self.strings,
             map: self.map,
+            resource_imports: self.resource_imports,
             allow_shadowing: self.allow_shadowing,
             _marker: self._marker,
         }
@@ -444,13 +497,19 @@ impl<T> LinkerInstance<'_, T> {
         name: &str,
         ty: ResourceType,
         dtor: impl Fn(StoreContextMut<'_, T>, u32) -> Result<()> + Send + Sync + 'static,
-    ) -> Result<()> {
+    ) -> Result<ResourceImportIndex> {
         let name = self.strings.intern(name);
         let dtor = Arc::new(crate::func::HostFunc::wrap(
             &self.engine,
             move |mut cx: crate::Caller<'_, T>, param: u32| dtor(cx.as_context_mut(), param),
         ));
-        self.insert(name, Definition::Resource(ty, dtor))
+        let idx = ResourceImportIndex::new(*self.resource_imports);
+        *self.resource_imports = self
+            .resource_imports
+            .checked_add(1)
+            .context("resource import count would overflow")?;
+        self.insert(name, Definition::Resource(idx, ty, dtor))?;
+        Ok(idx)
     }
 
     /// Defines a nested instance within this instance.

diff --git a/crates/wasmtime/src/component/matching.rs b/crates/wasmtime/src/component/matching.rs
@@ -53,7 +53,7 @@ impl TypeChecker<'_> {
             TypeDef::Resource(i) => {
                 let i = self.types[i].ty;
                 let actual = match actual {
-                    Some(Definition::Resource(actual, _dtor)) => actual,
+                    Some(Definition::Resource(_idx, actual, _dtor)) => actual,
 
                     // If a resource is imported yet nothing was supplied then
                     // that's only successful if the resource has itself

diff --git a/crates/wasmtime/src/component/mod.rs b/crates/wasmtime/src/component/mod.rs
@@ -21,7 +21,7 @@ pub use self::func::{
     ComponentNamedList, ComponentType, Func, Lift, Lower, TypedFunc, WasmList, WasmStr,
 };
 pub use self::instance::{ExportInstance, Exports, Instance, InstancePre};
-pub use self::linker::{Linker, LinkerInstance};
+pub use self::linker::{Linker, LinkerInstance, ResourceImportIndex};
 pub use self::resource_table::{ResourceTable, ResourceTableError};
 pub use self::resources::{Resource, ResourceAny};
 pub use self::types::{ResourceType, Type};

diff --git a/crates/wasmtime/src/component/resources.rs b/crates/wasmtime/src/component/resources.rs
@@ -1,9 +1,11 @@
 use crate::component::func::{bad_type_info, desc, LiftContext, LowerContext};
+use crate::component::instance::RuntimeImport;
+use crate::component::linker::ResourceImportIndex;
 use crate::component::matching::InstanceType;
-use crate::component::{ComponentType, Lift, Lower};
+use crate::component::{ComponentType, InstancePre, Lift, Lower};
 use crate::store::{StoreId, StoreOpaque};
 use crate::{AsContextMut, StoreContextMut, Trap};
-use anyhow::{bail, Result};
+use anyhow::{bail, ensure, Context, Result};
 use std::any::TypeId;
 use std::fmt;
 use std::marker;
@@ -394,6 +396,16 @@ where
             _marker: marker::PhantomData,
         })
     }
+
+    /// See [`ResourceAny::try_from_resource`]
+    pub fn try_into_resource_any<U>(
+        self,
+        store: impl AsContextMut,
+        instance: &InstancePre<U>,
+        idx: ResourceImportIndex,
+    ) -> Result<ResourceAny> {
+        ResourceAny::try_from_resource(self, store, instance, idx)
+    }
 }
 
 unsafe impl<T: 'static> ComponentType for Resource<T> {
@@ -470,9 +482,11 @@ impl<T> fmt::Debug for Resource<T> {
 /// This type is similar to [`Resource`] except that it can be used to represent
 /// any resource, either host or guest. This type cannot be directly constructed
 /// and is only available if the guest returns it to the host (e.g. a function
-/// returning a guest-defined resource). This type also does not carry a static
-/// type parameter `T` for example and does not have as much information about
-/// its type. This means that it's possible to get runtime type-errors when
+/// returning a guest-defined resource) or by a conversion from [`Resource`] via
+/// [`ResourceAny::try_from_resource`].
+/// This type also does not carry a static type parameter `T` for example and
+/// does not have as much information about its type.
+/// This means that it's possible to get runtime type-errors when
 /// using this type because it cannot statically prevent mismatching resource
 /// types.
 ///
@@ -501,6 +515,59 @@ struct OwnState {
 }
 
 impl ResourceAny {
+    /// Attempts to convert an imported [`Resource`] into [`ResourceAny`].
+    /// `idx` is the [`ResourceImportIndex`] returned by [`Linker::resource`].
+    ///
+    /// [`Linker::resource`]: crate::component::LinkerInstance::resource
+    pub fn try_from_resource<T: 'static, U>(
+        Resource { rep, state, .. }: Resource<T>,
+        mut store: impl AsContextMut,
+        instance_pre: &InstancePre<U>,
+        idx: ResourceImportIndex,
+    ) -> Result<Self> {
+        let store = store.as_context_mut();
+        let import = instance_pre
+            .resource_import(idx)
+            .context("import not found")?;
+        let RuntimeImport::Resource {
+            ty, dtor_funcref, ..
+        } = import
+        else {
+            bail!("import is not a resource")
+        };
+        ensure!(*ty == ResourceType::host::<T>(), "resource type mismatch");
+
+        let mut tables = host_resource_tables(store.0);
+        let (idx, own_state) = match state.load(Relaxed) {
+            BORROW => (tables.resource_lower_borrow(None, rep), None),
+            NOT_IN_TABLE => {
+                let idx = tables.resource_lower_own(None, rep);
+                (
+                    idx,
+                    Some(OwnState {
+                        dtor: Some(dtor_funcref.into()),
+                        flags: None,
+                        store: store.0.id(),
+                    }),
+                )
+            }
+            TAKEN => bail!("host resource already consumed"),
+            idx => (
+                idx,
+                Some(OwnState {
+                    dtor: Some(dtor_funcref.into()),
+                    flags: None,
+                    store: store.0.id(),
+                }),
+            ),
+        };
+        Ok(Self {
+            idx,
+            ty: *ty,
+            own_state,
+        })
+    }
+
     /// Returns the corresponding type associated with this resource, either a
     /// host-defined type or a guest-defined type.
     ///