[wit-component] fix task.return name matching (#1979)

`wit-bingen` emits world-level `task.return` imports using a module name of
"[export]$root", but `wit-component` was expecting just "$root".  This updates
`wit-component` to match what `wit-bindgen` emits.

Signed-off-by: Joel Dice <[email protected]>

crates/wit-component/src/validation.rs

@@ -621,10 +621,6 @@ impl ImportMap {
                 });
             }
 
-            if let Some(import) = async_import_for_export(None)? {
-                return Ok(import);
-            }
-
             let key = WorldKey::Name(name.to_string());
             if let Some(WorldItem::Function(func)) = world.imports.get(&key) {
                 validate_func(resolve, ty, func, abi)?;
@@ -646,6 +642,15 @@ impl ImportMap {
             }
         }
 
+        if matches!(
+            module.strip_prefix(names.import_exported_intrinsic_prefix()),
+            Some(module) if module == names.import_root()
+        ) {
+            if let Some(import) = async_import_for_export(None)? {
+                return Ok(import);
+            }
+        }
+
         let interface = match module.strip_prefix(names.import_non_root_prefix()) {
             Some(name) => name,
             None => bail!("unknown or invalid component model import syntax"),

crates/wit-component/tests/components/async-builtins/component.wat

@@ -7,7 +7,7 @@
     (type (;4;) (func (param i32 i32) (result i32)))
     (type (;5;) (func (param i32 i32 i32 i32) (result i32)))
     (import "$root" "[task-backpressure]" (func (;0;) (type 0)))
-    (import "$root" "[task-return]foo" (func (;1;) (type 1)))
+    (import "[export]$root" "[task-return]foo" (func (;1;) (type 1)))
     (import "[export]foo:foo/bar" "[task-return]foo" (func (;2;) (type 1)))
     (import "$root" "[task-wait]" (func (;3;) (type 2)))
     (import "$root" "[task-poll]" (func (;4;) (type 2)))
@@ -87,60 +87,63 @@
   )
   (core instance (;0;) (instantiate 1))
   (core func (;0;) (canon task.backpressure))
-  (core type (;0;) (func (param i32 i32)))
-  (core func (;1;) (canon task.return 0))
-  (alias core export 0 "0" (core func (;2;)))
-  (alias core export 0 "1" (core func (;3;)))
-  (core func (;4;) (canon task.yield))
-  (core func (;5;) (canon subtask.drop))
-  (alias core export 0 "2" (core func (;6;)))
-  (alias core export 0 "3" (core func (;7;)))
-  (core func (;8;) (canon error-context.drop))
+  (alias core export 0 "0" (core func (;1;)))
+  (alias core export 0 "1" (core func (;2;)))
+  (core func (;3;) (canon task.yield))
+  (core func (;4;) (canon subtask.drop))
+  (alias core export 0 "2" (core func (;5;)))
+  (alias core export 0 "3" (core func (;6;)))
+  (core func (;7;) (canon error-context.drop))
   (core instance (;1;)
     (export "[task-backpressure]" (func 0))
-    (export "[task-return]foo" (func 1))
-    (export "[task-wait]" (func 2))
-    (export "[task-poll]" (func 3))
-    (export "[task-yield]" (func 4))
-    (export "[subtask-drop]" (func 5))
-    (export "[error-context-new;encoding=utf8]" (func 6))
-    (export "[error-context-debug-message;encoding=utf8;realloc=cabi_realloc]" (func 7))
-    (export "[error-context-drop]" (func 8))
+    (export "[task-wait]" (func 1))
+    (export "[task-poll]" (func 2))
+    (export "[task-yield]" (func 3))
+    (export "[subtask-drop]" (func 4))
+    (export "[error-context-new;encoding=utf8]" (func 5))
+    (export "[error-context-debug-message;encoding=utf8;realloc=cabi_realloc]" (func 6))
+    (export "[error-context-drop]" (func 7))
+  )
+  (core type (;0;) (func (param i32 i32)))
+  (core func (;8;) (canon task.return 0))
+  (core instance (;2;)
+    (export "[task-return]foo" (func 8))
   )
   (core type (;1;) (func (param i32 i32)))
   (core func (;9;) (canon task.return 1))
-  (core instance (;2;)
+  (core instance (;3;)
     (export "[task-return]foo" (func 9))
   )
-  (core instance (;3;) (instantiate 0
+  (core instance (;4;) (instantiate 0
       (with "$root" (instance 1))
-      (with "[export]foo:foo/bar" (instance 2))
+      (with "[export]$root" (instance 2))
+      (with "[export]foo:foo/bar" (instance 3))
     )
   )
-  (alias core export 3 "memory" (core memory (;0;)))
+  (alias core export 4 "memory" (core memory (;0;)))
   (alias core export 0 "$imports" (core table (;0;)))
   (core func (;10;) (canon task.wait (memory 0)))
   (core func (;11;) (canon task.poll (memory 0)))
   (core func (;12;) (canon error-context.new (memory 0) string-encoding=utf8))
-  (alias core export 3 "cabi_realloc" (core func (;13;)))
+  (alias core export 4 "cabi_realloc" (core func (;13;)))
   (core func (;14;) (canon error-context.debug-message (memory 0) (realloc 13) string-encoding=utf8))
-  (core instance (;4;)
+  (core instance (;5;)
     (export "$imports" (table 0))
     (export "0" (func 10))
     (export "1" (func 11))
     (export "2" (func 12))
     (export "3" (func 14))
   )
-  (core instance (;5;) (instantiate 2
-      (with "" (instance 4))
+  (core instance (;6;) (instantiate 2
+      (with "" (instance 5))
     )
   )
   (type (;0;) (func (param "s" string) (result string)))
-  (alias core export 3 "[async-stackful]foo" (core func (;15;)))
+  (alias core export 4 "[async-stackful]foo" (core func (;15;)))
   (func (;0;) (type 0) (canon lift (core func 15) (memory 0) (realloc 13) string-encoding=utf8 async))
   (export (;1;) "foo" (func 0))
   (type (;1;) (func (param "s" string) (result string)))
-  (alias core export 3 "[async-stackful]foo:foo/bar#foo" (core func (;16;)))
+  (alias core export 4 "[async-stackful]foo:foo/bar#foo" (core func (;16;)))
   (func (;2;) (type 1) (canon lift (core func 16) (memory 0) (realloc 13) string-encoding=utf8 async))
   (component (;0;)
     (type (;0;) (func (param "s" string) (result string)))

crates/wit-component/tests/components/async-builtins/module.wat

@@ -1,6 +1,6 @@
 (module
   (import "$root" "[task-backpressure]" (func (param i32)))
-  (import "$root" "[task-return]foo" (func (param i32 i32)))
+  (import "[export]$root" "[task-return]foo" (func (param i32 i32)))
   (import "[export]foo:foo/bar" "[task-return]foo" (func (param i32 i32)))  
   (import "$root" "[task-wait]" (func (param i32) (result i32)))
   (import "$root" "[task-poll]" (func (param i32) (result i32)))