Index stream_payloads and future_payloads maps with Option<Type>

crates/rust/src/bindgen.rs

@@ -1,6 +1,6 @@
 use crate::{
-    ConstructorReturnType, Identifier, InterfaceGenerator, RustFlagsRepr,
-    classify_constructor_return_type, int_repr, to_rust_ident,
+    ConstructorReturnType, InterfaceGenerator, RustFlagsRepr, classify_constructor_return_type,
+    int_repr, to_rust_ident,
 };
 use heck::*;
 use std::fmt::Write as _;
@@ -427,19 +427,8 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             Instruction::FutureLift { payload, .. } => {
                 let async_support = self.r#gen.r#gen.async_support_path();
                 let op = &operands[0];
-                let name = payload
-                    .as_ref()
-                    .map(|ty| {
-                        self.r#gen
-                            .type_name_owned_with_id(ty, Identifier::StreamOrFuturePayload)
-                    })
-                    .unwrap_or_else(|| "()".into());
-                let ordinal = self
-                    .r#gen
-                    .r#gen
-                    .future_payloads
-                    .get_index_of(&name)
-                    .unwrap();
+                let key = payload.as_ref().map(|ty| self.r#gen.canonical_type_key(ty));
+                let ordinal = self.r#gen.r#gen.future_payloads.get_index_of(&key).unwrap();
                 let path = self.r#gen.path_to_root();
                 results.push(format!(
                     "{async_support}::FutureReader::new\
@@ -455,19 +444,8 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             Instruction::StreamLift { payload, .. } => {
                 let async_support = self.r#gen.r#gen.async_support_path();
                 let op = &operands[0];
-                let name = payload
-                    .as_ref()
-                    .map(|ty| {
-                        self.r#gen
-                            .type_name_owned_with_id(ty, Identifier::StreamOrFuturePayload)
-                    })
-                    .unwrap_or_else(|| "()".into());
-                let ordinal = self
-                    .r#gen
-                    .r#gen
-                    .stream_payloads
-                    .get_index_of(&name)
-                    .unwrap();
+                let key = payload.as_ref().map(|ty| self.r#gen.canonical_type_key(ty));
+                let ordinal = self.r#gen.r#gen.stream_payloads.get_index_of(&key).unwrap();
                 let path = self.r#gen.path_to_root();
                 results.push(format!(
                     "{async_support}::StreamReader::new\

crates/rust/src/interface.rs

@@ -528,12 +528,13 @@ macro_rules! {macro_name} {{
         } else {
             "()".into()
         };
+        let key = payload_type.map(|ty| self.canonical_type_key(ty));
         let map = match payload_for {
             PayloadFor::Future => &mut self.r#gen.future_payloads,
             PayloadFor::Stream => &mut self.r#gen.stream_payloads,
         };
 
-        if map.contains_key(&name) {
+        if map.contains_key(&key) {
             return;
         }
         let ordinal = map.len();
@@ -682,7 +683,7 @@ pub mod vtable{ordinal} {{
             PayloadFor::Future => &mut self.r#gen.future_payloads,
             PayloadFor::Stream => &mut self.r#gen.stream_payloads,
         };
-        map.insert(name, code);
+        map.insert(key, code);
     }
 
     fn generate_guest_import(&mut self, func: &Function, interface: Option<&WorldKey>) {
@@ -1700,11 +1701,16 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
         }
     }
 
-    pub(crate) fn type_name_owned_with_id(&mut self, ty: &Type, id: Identifier<'i>) -> String {
-        let old_identifier = mem::replace(&mut self.identifier, id);
-        let name = self.type_name_owned(ty);
-        self.identifier = old_identifier;
-        name
+    /// Returns the canonical representation of the given Type for the purposes
+    /// of deduplicating StreamPayload and FuturePayload implementations.
+    ///
+    /// If the given Type is a TypeId, then dealias it. Otherwise, return the
+    /// primitive type passed in.
+    pub(crate) fn canonical_type_key(&self, t: &Type) -> Type {
+        match t {
+            Type::Id(id) => Type::Id(dealias(self.resolve, *id)),
+            _ => *t,
+        }
     }
 
     fn type_name_owned(&mut self, ty: &Type) -> String {

crates/rust/src/lib.rs

@@ -51,8 +51,8 @@ pub struct RustWasm {
     /// Maps wit interface and type names to their Rust identifiers
     with: GenerationConfiguration,
 
-    future_payloads: IndexMap<String, String>,
-    stream_payloads: IndexMap<String, String>,
+    future_payloads: IndexMap<Option<Type>, String>,
+    stream_payloads: IndexMap<Option<Type>, String>,
 }
 
 #[derive(Default)]