wasmtime-cli: component run's --invoke uses ItemName for search

Cargo.toml

@@ -566,6 +566,8 @@ pooling-allocator = ["wasmtime/pooling-allocator", "wasmtime-cli-flags/pooling-a
 backtrace = ["wasmtime/backtrace"]
 component-model = [
   "wasmtime/component-model",
+  "wasmtime/wit-parser",
+  "wasmtime/anyhow",
   "wasmtime-wast?/component-model",
   "wasmtime-cli-flags/component-model",
   "wasmtime-wizer?/component-model",

src/commands/run.rs

@@ -796,29 +796,20 @@ impl RunCommand {
                 )
         })?;
 
-        let name = untyped_call.name();
-        let matches =
-            Self::search_component_funcs(store.engine(), component.component_type(), name);
-        let (names, func_type) = match matches.len() {
-            0 => bail!("No exported func named `{name}` in component."),
-            1 => &matches[0],
-            _ => bail!(
-                "Multiple exports named `{name}`: {matches:?}. FIXME: support some way to disambiguate names"
-            ),
-        };
+        let name = untyped_call.item_name().map_err(|e| {
+            wasmtime::Error::from_anyhow(e).context(format!(
+                "parsing `{}` as a wit item name",
+                untyped_call.name()
+            ))
+        })?;
 
-        let param_types = WasmFunc::params(func_type).collect::<Vec<_>>();
+        let (export, func_type) = Self::search_component_funcs(store, &component, &name)?;
+
+        let param_types = WasmFunc::params(&func_type).collect::<Vec<_>>();
         let params = untyped_call
             .to_wasm_params(&param_types)
             .with_context(|| format!("while interpreting parameters in invoke \"{invoke}\""))?;
 
-        let export = names
-            .iter()
-            .fold(None, |instance, name| {
-                component.get_export_index(instance.as_ref(), name)
-            })
-            .expect("export has at least one name");
-
         let instance = linker.instantiate_async(&mut *store, component).await?;
 
         let func = instance
@@ -944,47 +935,90 @@ impl RunCommand {
 
     #[cfg(feature = "component-model")]
     fn search_component_funcs(
-        engine: &Engine,
-        component: wasmtime::component::types::Component,
-        name: &str,
-    ) -> Vec<(Vec<String>, wasmtime::component::types::ComponentFunc)> {
+        store: &mut Store<Host>,
+        component: &wasmtime::component::Component,
+        item_name: &wasmtime::component::wit_parser::ItemName,
+    ) -> Result<(
+        wasmtime::component::ComponentExportIndex,
+        wasmtime::component::types::ComponentFunc,
+    )> {
         use wasmtime::component::types::ComponentItem as CItem;
-        fn collect_exports(
-            engine: &Engine,
-            item: CItem,
-            basename: Vec<String>,
-        ) -> Vec<(Vec<String>, CItem)> {
-            match item {
-                CItem::Component(c) => c
-                    .exports(engine)
-                    .flat_map(move |(name, item)| {
-                        let mut names = basename.clone();
-                        names.push(name.to_string());
-                        collect_exports(engine, item.ty, names)
-                    })
-                    .collect::<Vec<_>>(),
-                CItem::ComponentInstance(c) => c
-                    .exports(engine)
-                    .flat_map(move |(name, item)| {
-                        let mut names = basename.clone();
-                        names.push(name.to_string());
-                        collect_exports(engine, item.ty, names)
-                    })
-                    .collect::<Vec<_>>(),
-                _ => vec![(basename, item)],
-            }
+        // Start by looking up the item name directly.
+        // Only match this as the search if it provides a function - it may
+        // provide an instance of the same name, in which case we want the
+        // below to search through it.
+        match component.get_export(None, item_name) {
+            Some((CItem::ComponentFunc(func), index)) => return Ok((index, func.clone())),
+            _ => {}
+        }
+        if item_name.interface.is_some() || item_name.package.is_some() {
+            // If the item name specified a package or interface, and the
+            // ItemName based lookup failed to find it, we do not consider
+            // that it may be exported under an instance and terminate the
+            // search immediately:
+            bail!("No exported func named `{item_name}` in component.")
         }
+        // If the item name does not specify a package or interface, and it
+        // wasn't found in the root of the component above, then we search all
+        // instance exports for a function by that name.
+        let needle = item_name.to_string();
+        let mut search = component
+            .component_type()
+            .exports(store.engine())
+            .filter_map(|(instname, item)| match item.ty {
+                CItem::ComponentInstance(inst) => {
+                    inst.exports(store.engine())
+                        .find_map(|(leafname, item)| match item.ty {
+                            CItem::ComponentFunc(func) => {
+                                if leafname == needle {
+                                    // The type::Component traversal
+                                    // didn't give us an export index, get
+                                    // that now:
+                                    let (_item, inst_index) = component
+                                        .get_export(None, instname)
+                                        .expect("found exported component instance");
+                                    let (_item, index) = component
+                                        .get_export(Some(&inst_index), leafname)
+                                        .expect("found func");
+                                    Some((index, func.clone(), instname.to_string()))
+                                } else {
+                                    None
+                                }
+                            }
+                            _ => None,
+                        })
+                }
 
-        collect_exports(engine, CItem::Component(component), Vec::new())
-            .into_iter()
-            .filter_map(|(names, item)| {
-                let CItem::ComponentFunc(func) = item else {
-                    return None;
-                };
-                let func_name = names.last().expect("at least one name");
-                (func_name == name).then_some((names, func))
+                _ => None,
             })
-            .collect()
+            .collect::<Vec<_>>();
+
+        match search.len() {
+            0 => bail!("No exported func named `{needle}` in component."),
+            1 => {
+                let (index, func, _instname) = search.pop().unwrap();
+                Ok((index, func))
+            }
+            _ => {
+                let candidates = search
+                    .into_iter()
+                    .map(|(_index, _func, instname)| {
+                        // Manipulate as an ItemName to get package version
+                        // correct in the output
+                        let mut itemname: wasmtime::component::wit_parser::ItemName =
+                            instname.parse().unwrap();
+                        // Push the function name onto the itemname:
+                        itemname.interface = Some(itemname.name.clone());
+                        itemname.name = needle.to_string();
+                        format!("`{itemname}`")
+                    })
+                    .collect::<Vec<_>>();
+                bail!(
+                    "Multiple instances contained funcs named `{needle}`, retry with a more specific name: {}",
+                    candidates.join(", ")
+                )
+            }
+        }
     }
 
     async fn invoke_func(&self, store: &mut Store<Host>, func: Func) -> Result<()> {

tests/all/cli_tests.rs

@@ -805,6 +805,55 @@ fn component_enabled_by_default() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn component_invoke_multiple_run_exports() -> Result<()> {
+    let path = "tests/all/cli_tests/component-multiple-runs.wat";
+    let wasm = build_wasm(path)?;
+
+    // demonstrate run --invoke can give a useful error message when
+    // there are multiple interfaces that export the function name specified
+    let output = get_wasmtime_command()?
+        .arg("run")
+        .arg("-Wcomponent-model")
+        .arg("-Ccache=n")
+        .arg("--invoke")
+        .arg("run()")
+        .arg(wasm.path())
+        .output()?;
+    assert!(
+        !output.status.success(),
+        "should fail because run() is ambigious for this component"
+    );
+    let stderr = String::from_utf8_lossy(&output.stderr);
+    assert!(stderr.contains("Multiple instances contained funcs named `run`, retry with a more specific name: `wasi:cli/run.run@0.2.0`, `some:other/one.run`"));
+
+    // test program cli run gives `ok`:
+    let output = get_wasmtime_command()?
+        .arg("run")
+        .arg("-Wcomponent-model")
+        .arg("-Ccache=n")
+        .arg("--invoke")
+        .arg("wasi:cli/run.run@0.2.0()")
+        .arg(wasm.path())
+        .output()?;
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert_eq!(stdout, "ok\n");
+
+    // test program other run gives `err`:
+    let output = get_wasmtime_command()?
+        .arg("run")
+        .arg("-Wcomponent-model")
+        .arg("-Ccache=n")
+        .arg("--invoke")
+        .arg("some:other/one.run()")
+        .arg(wasm.path())
+        .output()?;
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    assert_eq!(stdout, "err\n");
+
+    Ok(())
+}
+
 // If the text format is invalid then the filename should be mentioned in the
 // error message.
 #[test]

tests/all/cli_tests/component-multiple-runs.wat

@@ -0,0 +1,22 @@
+;; This component has two different interfaces each exporting a function
+;; called "run", which behave distinctly. wasi:cli/run's run returns ok,
+;; some:other's run returns error.
+(component
+  (core module $m
+    (func (export "run") (result i32)
+      i32.const 0)
+    (func (export "run2") (result i32)
+      i32.const 1)
+  )
+  (core instance $i (instantiate $m))
+  (func $run (result (result))
+    (canon lift (core func $i "run")))
+  (func $run2 (result (result))
+    (canon lift (core func $i "run2")))
+
+  (instance (export (interface "wasi:cli/run@0.2.0"))
+    (export "run" (func $run)))
+
+  (instance (export (interface "some:other/one"))
+    (export "run" (func $run2)))
+)