From d7d96d1d5ac4f8ab2cd3d46db211098002be9632 Mon Sep 17 00:00:00 2001
From: Roman Volosatovs <rvolosatovs@riseup.net>
Date: Thu, 21 Dec 2023 17:30:44 +0100
Subject: [PATCH] feat: support `RuntimeImportIndex` lookups by string paths

Signed-off-by: Roman Volosatovs <rvolosatovs@riseup.net>
---
 crates/wasmtime/src/component/instance.rs  | 26 ++++++++++++++++------
 crates/wasmtime/src/component/resources.rs | 15 ++++++-------
 tests/all/component_model/resources.rs     |  2 ++
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/crates/wasmtime/src/component/instance.rs b/crates/wasmtime/src/component/instance.rs
index 63f70dbab4f..c4cf49e03cf 100644
--- a/crates/wasmtime/src/component/instance.rs
+++ b/crates/wasmtime/src/component/instance.rs
@@ -558,15 +558,27 @@ impl<T> InstancePre<T> {
         }
     }
 
-    pub(crate) fn resource_import_index(
-        &self,
-        path: ResourceImportIndex,
-    ) -> Option<RuntimeImportIndex> {
-        *self.resource_imports.get(path)?
+    /// Returns [`RuntimeImportIndex`] associated with this resource.
+    /// `idx` is the [`ResourceImportIndex`] returned by [`Linker::resource`].
+    ///
+    /// [`Linker::resource`]: crate::component::LinkerInstance::resource
+    pub fn resource_import_index(&self, idx: ResourceImportIndex) -> Option<RuntimeImportIndex> {
+        *self.resource_imports.get(idx)?
+    }
+
+    /// Returns [`RuntimeImportIndex`] associated with this `path` within a `root`.
+    pub fn path_import_index(&self, root: &str, path: &[&str]) -> Option<RuntimeImportIndex> {
+        let env_component = self.component().env_component();
+        env_component
+            .imports
+            .iter()
+            .find_map(|(idx, (import, import_path))| {
+                let (root_name, _) = env_component.import_types.get(*import)?;
+                (root_name == root && import_path == path).then_some(idx)
+            })
     }
 
-    pub(crate) fn resource_import(&self, path: ResourceImportIndex) -> Option<&RuntimeImport> {
-        let idx = self.resource_import_index(path)?;
+    pub(crate) fn runtime_import(&self, idx: RuntimeImportIndex) -> Option<&RuntimeImport> {
         self.imports.get(idx)
     }
 
diff --git a/crates/wasmtime/src/component/resources.rs b/crates/wasmtime/src/component/resources.rs
index 624514d1ffe..a58e9c98391 100644
--- a/crates/wasmtime/src/component/resources.rs
+++ b/crates/wasmtime/src/component/resources.rs
@@ -1,6 +1,5 @@
 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, InstancePre, Lift, Lower};
 use crate::store::{StoreId, StoreOpaque};
@@ -11,7 +10,9 @@ use std::fmt;
 use std::marker;
 use std::mem::MaybeUninit;
 use std::sync::atomic::{AtomicU32, Ordering::Relaxed};
-use wasmtime_environ::component::{CanonicalAbiInfo, DefinedResourceIndex, InterfaceType};
+use wasmtime_environ::component::{
+    CanonicalAbiInfo, DefinedResourceIndex, InterfaceType, RuntimeImportIndex,
+};
 use wasmtime_runtime::component::{ComponentInstance, InstanceFlags, ResourceTables};
 use wasmtime_runtime::{SendSyncPtr, VMFuncRef, ValRaw};
 
@@ -428,7 +429,7 @@ where
         self,
         store: impl AsContextMut,
         instance: &InstancePre<U>,
-        idx: ResourceImportIndex,
+        idx: RuntimeImportIndex,
     ) -> Result<ResourceAny> {
         ResourceAny::try_from_resource(self, store, instance, idx)
     }
@@ -542,18 +543,16 @@ 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
+    /// `idx` is the [`RuntimeImportIndex`] associated with this resource.
     pub fn try_from_resource<T: 'static, U>(
         Resource { rep, state, .. }: Resource<T>,
         mut store: impl AsContextMut,
         instance_pre: &InstancePre<U>,
-        idx: ResourceImportIndex,
+        idx: RuntimeImportIndex,
     ) -> Result<Self> {
         let store = store.as_context_mut();
         let import = instance_pre
-            .resource_import(idx)
+            .runtime_import(idx)
             .context("import not found")?;
         let RuntimeImport::Resource {
             ty, dtor_funcref, ..
diff --git a/tests/all/component_model/resources.rs b/tests/all/component_model/resources.rs
index 448eceecf7a..18ff22a52e5 100644
--- a/tests/all/component_model/resources.rs
+++ b/tests/all/component_model/resources.rs
@@ -562,6 +562,8 @@ fn dynamic_val() -> Result<()> {
         .root()
         .resource("t1", ResourceType::host::<MyType>(), |_, _| Ok(()))?;
     let i_pre = linker.instantiate_pre(&c)?;
+    let idx = i_pre.resource_import_index(idx).unwrap();
+    assert_eq!(i_pre.path_import_index("t1", &[]), Some(idx));
     let i = i_pre.instantiate(&mut store)?;
 
     let a = i.get_func(&mut store, "a").unwrap();