Support special import names for unit futures/streams (#2390)

* Support special import names for unit futures/streams

This commit enhances `wit-component`'s import detection to support
"unit" `future` and `stream` types which don't have a payload in a
special way. Previously it was required that these types were mentioned
in some WIT somewhere to get connected to a function, but now it's
possible to import these intrinsics directly even if the type isn't
mentioned in WIT.

The motivation for this change is to support guest language async
runtimes which are world-agnostic to still operate on unit
futures/streams as necessary without binding them to any particular WIT
import or export. For example it's expected that unit streams are a
primitive for inter-task communication in the component model.

* Update crates/wit-component/src/validation.rs

Co-authored-by: Joel Dice <[email protected]>

---------

Co-authored-by: Joel Dice <[email protected]>

Author: alexcrichton

crates/wit-component/src/encoding.rs

@@ -73,7 +73,9 @@
 
 use crate::StringEncoding;
 use crate::metadata::{self, Bindgen, ModuleMetadata};
-use crate::validation::{Export, ExportMap, Import, ImportInstance, ImportMap, PayloadInfo};
+use crate::validation::{
+    Export, ExportMap, Import, ImportInstance, ImportMap, PayloadInfo, PayloadType,
+};
 use anyhow::{Context, Result, anyhow, bail};
 use indexmap::{IndexMap, IndexSet};
 use std::borrow::Cow;
@@ -1337,20 +1339,21 @@ impl<'a> EncodingState<'a> {
                     let metadata = self.info.module_metadata_for(*for_module);
                     let exports = self.info.exports_for(*for_module);
                     let instance_index = self.instance_for(*for_module);
-                    let (encoding, realloc) = if info.imported {
-                        (
-                            metadata
-                                .import_encodings
-                                .get(resolve, &info.key, &info.function),
-                            exports.import_realloc_for(info.interface, &info.function),
-                        )
-                    } else {
-                        (
-                            metadata
-                                .export_encodings
-                                .get(resolve, &info.key, &info.function),
-                            exports.export_realloc_for(&info.key, &info.function),
-                        )
+                    let (encoding, realloc) = match &info.ty {
+                        PayloadType::Type { function, .. } => {
+                            if info.imported {
+                                (
+                                    metadata.import_encodings.get(resolve, &info.key, function),
+                                    exports.import_realloc_for(info.interface, function),
+                                )
+                            } else {
+                                (
+                                    metadata.export_encodings.get(resolve, &info.key, function),
+                                    exports.export_realloc_for(&info.key, function),
+                                )
+                            }
+                        }
+                        PayloadType::UnitFuture | PayloadType::UnitStream => (None, None),
                     };
                     let encoding = encoding.unwrap_or(StringEncoding::UTF8);
                     let realloc_index = realloc.map(|name| {
@@ -1506,16 +1509,19 @@ impl<'a> EncodingState<'a> {
         // Finally though exports do use `info.interface`. Honestly I'm not
         // really entirely sure why. Fuzzing is happy though, and truly
         // everything must be ok if the fuzzers are happy, right?
-        let ComponentValType::Type(type_index) = if info.imported || info.interface.is_none() {
+        let mut encoder = if info.imported || info.interface.is_none() {
             self.root_import_type_encoder(None)
         } else {
             self.root_export_type_encoder(info.interface)
-        }
-        .encode_valtype(resolve, &Type::Id(info.ty))?
-        else {
-            unreachable!()
         };
-        Ok(type_index)
+        match info.ty {
+            PayloadType::Type { id, .. } => match encoder.encode_valtype(resolve, &Type::Id(id))? {
+                ComponentValType::Type(index) => Ok(index),
+                ComponentValType::Primitive(_) => unreachable!(),
+            },
+            PayloadType::UnitFuture => Ok(encoder.encode_unit_future()),
+            PayloadType::UnitStream => Ok(encoder.encode_unit_stream()),
+        }
     }
 
     /// This is a helper function that will declare any types necessary for

crates/wit-component/src/encoding/types.rs

@@ -20,6 +20,8 @@ pub struct TypeEncodingMaps<'a> {
     pub id_to_index: HashMap<TypeId, u32>,
     pub def_to_index: HashMap<&'a TypeDefKind, u32>,
     pub func_type_map: HashMap<FunctionKey<'a>, u32>,
+    pub unit_future: Option<u32>,
+    pub unit_stream: Option<u32>,
 }
 
 impl<'a> TypeEncodingMaps<'a> {
@@ -372,6 +374,26 @@ pub trait ValtypeEncoder<'a> {
         encoder.stream(ty);
         Ok(ComponentValType::Type(index))
     }
+
+    fn encode_unit_future(&mut self) -> u32 {
+        if let Some(index) = self.type_encoding_maps().unit_future {
+            return index;
+        }
+        let (index, encoder) = self.defined_type();
+        encoder.future(None);
+        self.type_encoding_maps().unit_future = Some(index);
+        index
+    }
+
+    fn encode_unit_stream(&mut self) -> u32 {
+        if let Some(index) = self.type_encoding_maps().unit_stream {
+            return index;
+        }
+        let (index, encoder) = self.defined_type();
+        encoder.stream(None);
+        self.type_encoding_maps().unit_stream = Some(index);
+        index
+    }
 }
 
 pub struct RootTypeEncoder<'state, 'a> {

crates/wit-component/src/encoding/world.rs

@@ -1,6 +1,6 @@
 use super::{Adapter, ComponentEncoder, LibraryInfo, RequiredOptions};
 use crate::validation::{
-    Import, ImportMap, ValidatedModule, validate_adapter_module, validate_module,
+    Import, ImportMap, PayloadType, ValidatedModule, validate_adapter_module, validate_module,
 };
 use anyhow::{Context, Result};
 use indexmap::{IndexMap, IndexSet};
@@ -396,7 +396,9 @@ impl<'a> ComponentWorld<'a> {
                 | Import::FutureCancelWrite { info, async_: _ }
                 | Import::FutureDropReadable(info)
                 | Import::FutureDropWritable(info) => {
-                    live.add_type_id(resolve, info.ty);
+                    if let PayloadType::Type { id, .. } = info.ty {
+                        live.add_type_id(resolve, id);
+                    }
                 }
 
                 // The `task.return` intrinsic needs to be able to refer to the

crates/wit-component/src/validation.rs

@@ -3,7 +3,7 @@ use crate::{ComponentEncoder, StringEncoding};
 use anyhow::{Context, Result, anyhow, bail};
 use indexmap::{IndexMap, IndexSet, map::Entry};
 use std::fmt;
-use std::hash::{Hash, Hasher};
+use std::hash::Hash;
 use std::mem;
 use wasm_encoder::ExportKind;
 use wasmparser::names::{ComponentName, ComponentNameKind};
@@ -135,16 +135,16 @@ pub enum ImportInstance {
 /// type appears, consider encoding them in the name mangling stream on an
 /// individual basis, similar to how we encode `error-context.*` built-in
 /// imports.
-#[derive(Debug, Eq, PartialEq, Clone)]
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
 pub struct PayloadInfo {
     /// The original, mangled import name used to import this built-in
     /// (currently used only for hashing and debugging).
     pub name: String,
     /// The resolved type id for the `stream` or `future` type of interest.
-    pub ty: TypeId,
-    /// The component-level function import or export where the type appeared as
-    /// a parameter or result type.
-    pub function: String,
+    ///
+    /// If `Unit{Future,Stream}` this means that it's a "unit" payload or has no associated
+    /// type being sent.
+    pub ty: PayloadType,
     /// The world key representing the import or export context of `function`.
     pub key: WorldKey,
     /// The interface that `function` was imported from or exported in, if any.
@@ -156,30 +156,37 @@ pub struct PayloadInfo {
     pub imported: bool,
 }
 
+/// The type of future/stream referenced by a `PayloadInfo`
+#[derive(Debug, Eq, PartialEq, Clone, Hash)]
+pub enum PayloadType {
+    /// This is a future or stream located in a `Resolve` where `id` points to
+    /// either of `TypeDefKind::{Future, Stream}`.
+    Type {
+        id: TypeId,
+        /// The component-level function import or export where the type
+        /// appeared as a parameter or result type.
+        function: String,
+    },
+    /// This is a `future` (no type)
+    UnitFuture,
+    /// This is a `stream` (no type)
+    UnitStream,
+}
+
 impl PayloadInfo {
     /// Returns the payload type that this future/stream type is using.
     pub fn payload(&self, resolve: &Resolve) -> Option<Type> {
-        match resolve.types[self.ty].kind {
+        let id = match self.ty {
+            PayloadType::Type { id, .. } => id,
+            PayloadType::UnitFuture | PayloadType::UnitStream => return None,
+        };
+        match resolve.types[id].kind {
             TypeDefKind::Future(payload) | TypeDefKind::Stream(payload) => payload,
             _ => unreachable!(),
         }
     }
 }
 
-impl Hash for PayloadInfo {
-    /// We derive `Hash` for this type by hand and exclude the `function` field
-    /// because (A) `Function` doesn't implement `Hash` and (B) the other fields
-    /// are sufficient to uniquely identify the type of interest, which function
-    /// it appeared in, and which parameter or return type we found it in.
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.name.hash(state);
-        self.ty.hash(state);
-        self.key.hash(state);
-        self.interface.hash(state);
-        self.imported.hash(state);
-    }
-}
-
 /// The different kinds of items that a module or an adapter can import.
 ///
 /// This is intended to be an exhaustive definition of what can be imported into
@@ -1994,30 +2001,48 @@ impl Legacy {
     ) -> Option<PayloadInfo> {
         // parse the `prefix` into `func_name` and `type_index`, bailing out
         // with `None` if anything doesn't match.
-        let (type_index, func_name) = prefixed_integer(name, prefix)?;
-        let type_index = type_index as usize;
-
-        // Double-check that `func_name` is indeed a function name within
-        // this interface/world. Then additionally double-check that
-        // `type_index` is indeed a valid index for this function's type
-        // signature.
-        let function = get_function(
-            lookup_context.resolve,
-            lookup_context.world,
-            func_name,
-            lookup_context.id,
-            lookup_context.import,
-        )
-        .ok()?;
-        let ty = *function
-            .find_futures_and_streams(lookup_context.resolve)
-            .get(type_index)?;
+        let (index_or_unit, func_name) = prefixed_intrinsic(name, prefix)?;
+        let ty = match index_or_unit {
+            "unit" => {
+                if name.starts_with("[future") {
+                    PayloadType::UnitFuture
+                } else if name.starts_with("[stream") {
+                    PayloadType::UnitStream
+                } else {
+                    unreachable!()
+                }
+            }
+            other => {
+                // Note that this is parsed as a `u32` to ensure that the
+                // integer parsing is the same across platforms regardless of
+                // the the width of `usize`.
+                let type_index = other.parse::<u32>().ok()? as usize;
+
+                // Double-check that `func_name` is indeed a function name within
+                // this interface/world. Then additionally double-check that
+                // `type_index` is indeed a valid index for this function's type
+                // signature.
+                let function = get_function(
+                    lookup_context.resolve,
+                    lookup_context.world,
+                    func_name,
+                    lookup_context.id,
+                    lookup_context.import,
+                )
+                .ok()?;
+                PayloadType::Type {
+                    id: *function
+                        .find_futures_and_streams(lookup_context.resolve)
+                        .get(type_index)?,
+                    function: function.name.clone(),
+                }
+            }
+        };
 
         // And if all that passes wrap up everything in a `PayloadInfo`.
         Some(PayloadInfo {
             name: name.to_string(),
             ty,
-            function: function.name.clone(),
             key: lookup_context
                 .key
                 .clone()
@@ -2631,14 +2656,20 @@ fn validate_func_sig(name: &str, expected: &FuncType, ty: &wasmparser::FuncType)
     Ok(())
 }
 
-/// Matches `name` as `[${prefix}N]...`, and if found returns `(N, "...")`
-fn prefixed_integer<'a>(name: &'a str, prefix: &str) -> Option<(u32, &'a str)> {
+/// Matches `name` as `[${prefix}S]...`, and if found returns `("S", "...")`
+fn prefixed_intrinsic<'a>(name: &'a str, prefix: &str) -> Option<(&'a str, &'a str)> {
     assert!(prefix.starts_with("["));
     assert!(prefix.ends_with("-"));
     let suffix = name.strip_prefix(prefix)?;
     let index = suffix.find(']')?;
     let rest = &suffix[index + 1..];
-    let n = suffix[..index].parse().ok()?;
+    Some((&suffix[..index], rest))
+}
+
+/// Matches `name` as `[${prefix}N]...`, and if found returns `(N, "...")`
+fn prefixed_integer<'a>(name: &'a str, prefix: &str) -> Option<(u32, &'a str)> {
+    let (suffix, rest) = prefixed_intrinsic(name, prefix)?;
+    let n = suffix.parse().ok()?;
     Some((n, rest))
 }