rust: Reimplement how stream/future payloads work (#1528)

* rust: Reimplement how stream/future payloads work

Previously stream/future payload were generated by collecting the set of
types used in `future` and `stream` types in a WIT, rendering them to a
string, deduplicating based on this string representation, and then
generating various impls-with-vtables. This stringification strategy
unfortunately falls down in a few situations such as:

* Type aliases in WIT render as two names in Rust, but they're using the
  same Rust type.
* Types with the same definition, but in multiple modules, will have two
  different paths in Rust but alias the same type.
* Primitives may be used directly in streams/futures but then
  additionally used as a WIT type alias.

In all of these situations it's effectively exposing how Rust requires
at most one-impl-per-type-definition but the stringification/deduping
was just a proxy for implementing this restriction and not a precise
calculation. Using the work from bytecodealliance/wasm-tools#2447 as
well as #1468 it's possible to do all of this without stringifying.
Specifically #1468, transitively enabled by
bytecodealliance/wasm-tools#2447, enables building a set of equal types
that the Rust generator knows will all alias the same type definition.
Using this it's possible to translate a payload to its "canonical
payload" representation ID-wise and perform hashing/deduplication based
on that. This in turn solves all of the issues above as well as previous
issues such as #1432 and #1433 without requiring the workaround in #1482.

The end result is that all of these various bugs should be fixed and the
Rust generator should be much more reliable about when exactly a trait
impl is emitted vs not.

Closes #1523
Closes #1524

* Try fixing CI

* Fix typo

* Fix rust CI

Author: alexcrichton

crates/core/src/types.rs

@@ -94,8 +94,27 @@ impl Types {
             }
         }
     }
-    pub fn collect_equal_types(&mut self, resolve: &Resolve) {
+
+    /// Populates the return value of [`Types::get_representative_type`] with
+    /// the `resolve` passed in.
+    ///
+    /// The `may_alias_another_type` closure is used to determine whether the
+    /// language's definition of the provided `TypeId` might possibly alias
+    /// some other type in a language. This is a language-specific deduction
+    /// which can also be affected by options to a binding generator. If a type
+    /// can't be aliased by anything else then it can't be considered equal to
+    /// anything else. Note that in this situations other types may still
+    /// be equal to it, such as aliased types at the WIT level (e.g. `type foo
+    /// = some-record`).
+    pub fn collect_equal_types(
+        &mut self,
+        resolve: &Resolve,
+        may_alias_another_type: &dyn Fn(TypeId) -> bool,
+    ) {
         for (i, (ty, _)) in resolve.types.iter().enumerate() {
+            if !may_alias_another_type(ty) {
+                continue;
+            }
             // TODO: we could define a hash function for TypeDefKind to prevent the inner loop.
             for (earlier, _) in resolve.types.iter().take(i) {
                 if self.equal_types.find(ty) == self.equal_types.find(earlier) {

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 _;
@@ -428,28 +428,13 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 let async_support = self.r#gen.r#gen.async_support_path();
                 let op = &operands[0];
                 let name = match payload {
-                    Some(Type::Id(type_id)) => {
-                        let dealiased_id = dealias(resolve, *type_id);
-                        self.r#gen.type_name_owned_with_id(
-                            &Type::Id(dealiased_id),
-                            Identifier::StreamOrFuturePayload,
-                        )
-                    }
-                    Some(ty) => self
-                        .r#gen
-                        .type_name_owned_with_id(ty, Identifier::StreamOrFuturePayload),
-                    None => "()".into(),
+                    Some(ty) => self.r#gen.type_name_owned(ty),
+                    None => "()".to_string(),
                 };
-                let ordinal = self
-                    .r#gen
-                    .r#gen
-                    .future_payloads
-                    .get_index_of(&name)
-                    .unwrap();
                 let path = self.r#gen.path_to_root();
                 results.push(format!(
                     "{async_support}::FutureReader::new\
-                     ({op} as u32, &{path}wit_future::vtable{ordinal}::VTABLE)"
+                     ({op} as u32, &<{name} as {path}wit_future::FuturePayload>::VTABLE)"
                 ))
             }
 
@@ -462,28 +447,13 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 let async_support = self.r#gen.r#gen.async_support_path();
                 let op = &operands[0];
                 let name = match payload {
-                    Some(Type::Id(type_id)) => {
-                        let dealiased_id = dealias(resolve, *type_id);
-                        self.r#gen.type_name_owned_with_id(
-                            &Type::Id(dealiased_id),
-                            Identifier::StreamOrFuturePayload,
-                        )
-                    }
-                    Some(ty) => self
-                        .r#gen
-                        .type_name_owned_with_id(ty, Identifier::StreamOrFuturePayload),
-                    None => "()".into(),
+                    Some(ty) => self.r#gen.type_name_owned(ty),
+                    None => "()".to_string(),
                 };
-                let ordinal = self
-                    .r#gen
-                    .r#gen
-                    .stream_payloads
-                    .get_index_of(&name)
-                    .unwrap();
                 let path = self.r#gen.path_to_root();
                 results.push(format!(
                     "{async_support}::StreamReader::new\
-                     ({op} as u32, &{path}wit_stream::vtable{ordinal}::VTABLE)"
+                     ({op} as u32, &<{name} as {path}wit_stream::StreamPayload>::VTABLE)"
                 ))
             }
 
@@ -1281,7 +1251,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 self.push_str(&format!(
                     " let array{tmp}: [_; {size}] = core::array::from_fn(|{index_var}| {{
                             let base = {base}.add({index_var} * {elemsize});
-                            {body} 
+                            {body}
                         }});"
                 ));
                 let result = format!("array{tmp}");

crates/rust/src/interface.rs

@@ -523,11 +523,30 @@ macro_rules! {macro_name} {{
         func_name: &str,
         payload_type: Option<&Type>,
     ) {
-        let name = match payload_type {
-            Some(Type::Id(type_id)) => {
-                let dealiased_id = dealias(self.resolve, *type_id);
-                self.type_name_owned(&Type::Id(dealiased_id))
+        let payload_type = match payload_type {
+            // Rust requires one-impl-per-type, so any `id` here is transformed
+            // into its canonical representation using
+            // `get_representative_type`. This ensures that type aliases, uses,
+            // etc, all get canonicalized to the exact same ID regardless of
+            // type structure.
+            //
+            // Note that `get_representative_type` maps ids-to-ids which is 95%
+            // of what we want, but this additionally goes one layer further to
+            // see if the final id is actually itself a typedef, which would
+            // always be to a primitive, and then uses the primitive type
+            // instead of the typedef to canonicalize with other streams/futures
+            // using the primitive type.
+            Some(Type::Id(id)) => {
+                let id = self.r#gen.types.get_representative_type(*id);
+                match self.resolve.types[id].kind {
+                    TypeDefKind::Type(t) => Some(t),
+                    _ => Some(Type::Id(id)),
+                }
             }
+            other => other.copied(),
+        };
+        let payload_type = payload_type.as_ref();
+        let name = match payload_type {
             Some(payload_type) => self.type_name_owned(payload_type),
             None => "()".into(),
         };
@@ -536,7 +555,7 @@ macro_rules! {macro_name} {{
             PayloadFor::Stream => &mut self.r#gen.stream_payloads,
         };
 
-        if map.contains_key(&name) {
+        if map.contains_key(&payload_type.copied()) {
             return;
         }
         let ordinal = map.len();
@@ -685,7 +704,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(payload_type.copied(), code);
     }
 
     fn generate_guest_import(&mut self, func: &Function, interface: Option<&WorldKey>) {
@@ -1713,14 +1732,7 @@ 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
-    }
-
-    fn type_name_owned(&mut self, ty: &Type) -> String {
+    pub fn type_name_owned(&mut self, ty: &Type) -> String {
         self.type_name(
             ty,
             TypeMode {
@@ -2579,7 +2591,7 @@ impl<'a> wit_bindgen_core::InterfaceGenerator<'a> for InterfaceGenerator<'a> {
 
     fn define_type(&mut self, name: &str, id: TypeId) {
         let equal = self.r#gen.types.get_representative_type(id);
-        if equal == id {
+        if !self.r#gen.opts.merge_structurally_equal_types() || equal == id {
             wit_bindgen_core::define_type(self, name, id)
         } else {
             let docs = &self.resolve.types[id].docs;

crates/rust/src/lib.rs

@@ -50,8 +50,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)]
@@ -1129,9 +1129,39 @@ impl WorldGenerator for RustWasm {
             uwriteln!(self.src_preamble, "//   * async: {opt}");
         }
         self.types.analyze(resolve);
-        if self.opts.merge_structurally_equal_types() {
-            self.types.collect_equal_types(resolve);
-        }
+        self.types.collect_equal_types(resolve, &|a| {
+            // If `--merge-structurally-equal-types` is enabled then any type
+            // anywhere can be generated as a type alias to anything else.
+            if self.opts.merge_structurally_equal_types() {
+                return true;
+            }
+
+            match resolve.types[a].kind {
+                // These types are all defined with `type Foo = ...` in Rust
+                // since Rust either has native representations or they live in
+                // libraries or similar.
+                TypeDefKind::Type(_)
+                | TypeDefKind::Handle(_)
+                | TypeDefKind::List(_)
+                | TypeDefKind::Tuple(_)
+                | TypeDefKind::Option(_)
+                | TypeDefKind::Result(_)
+                | TypeDefKind::Future(_)
+                | TypeDefKind::Stream(_)
+                | TypeDefKind::Map(..)
+                | TypeDefKind::FixedLengthList(..) => true,
+
+                // These types are all defined with fresh new types defined
+                // in generated bindings and thus can't alias some other
+                // existing type.
+                TypeDefKind::Record(_)
+                | TypeDefKind::Variant(_)
+                | TypeDefKind::Enum(_)
+                | TypeDefKind::Flags(_)
+                | TypeDefKind::Resource
+                | TypeDefKind::Unknown => false,
+            }
+        });
         self.world = Some(world);
 
         let world = &resolve.worlds[world];

crates/test/src/cpp.rs

@@ -41,22 +41,11 @@ impl LanguageMethods for Cpp {
 
     fn should_fail_verify(
         &self,
-        name: &str,
-        _config: &crate::config::WitConfig,
+        _name: &str,
+        config: &crate::config::WitConfig,
         _args: &[String],
     ) -> bool {
-        match name {
-            "async-trait-function.wit"
-            | "error-context.wit"
-            | "futures.wit"
-            | "import-export-future.wit"
-            | "import-export-stream.wit"
-            | "resources-with-futures.wit"
-            | "resources-with-streams.wit"
-            | "streams.wit"
-            | "async-resource-func.wit" => true,
-            _ => false,
-        }
+        config.async_
     }
 
     fn prepare(&self, runner: &mut Runner) -> anyhow::Result<()> {

crates/test/src/csharp.rs

@@ -49,6 +49,9 @@ impl LanguageMethods for Csharp {
                 | "resource-fallible-constructor.wit"
                 | "async-resource-func.wit"
                 | "import-export-stream.wit"
+                | "issue-1432.wit"
+                | "issue-1433.wit"
+                | "future-same-type-different-names.wit"
         )
     }
 

crates/test/src/rust.rs

@@ -63,8 +63,9 @@ impl LanguageMethods for Rust {
         // no_std doesn't currently work with async
         if config.async_
             && args.iter().any(|s| s == "--std-feature")
-            // Except this one actually _does_ work:
+            // Except these actually do work:
             && name != "async-trait-function.wit-no-std"
+            && name != "async-resource-func.wit-no-std"
         {
             return true;
         }

tests/codegen/async-resource-func.wit

@@ -1,3 +1,5 @@
+//@ async = true
+
 package foo:bar;
 
 world bindings {

tests/codegen/future-same-type-different-names.wit

@@ -0,0 +1,67 @@
+//@ async = true
+
+package a:b;
+
+interface i {
+  resource a;
+  f1: func(x: future<a>);
+
+  type b = a;
+  f2: func(x: future<b>);
+}
+
+interface i2 {
+  use i.{a};
+  f3: func(x: future<a>);
+
+  use i.{b};
+  f4: func(x: future<b>);
+}
+
+interface i3 {
+  f5: func(x: future<u8>);
+
+  type hello = u8;
+  f6: func(x: future<hello>);
+}
+
+interface i4 {
+  type a = result<u8, u32>;
+  f7: func(x: future<a>);
+
+  type b = result<u8, u32>;
+  f8: func(x: future<b>);
+}
+
+interface i5 {
+  type a = result<u8, u32>;
+  f9: func(x: future<a>);
+}
+
+interface i6 {
+  enum error-code {
+    io,
+    closed,
+  }
+}
+
+interface i7 {
+  use i6.{error-code};
+  f10: func() -> future<result<_, error-code>>;
+}
+
+interface i8 {
+  use i6.{error-code};
+  f11: func() -> future<result<_, error-code>>;
+}
+
+world foo {
+  import i;
+  import i2;
+  import i3;
+  import i4;
+  import i5;
+  import i6;
+  import i7;
+  import i8;
+}

tests/codegen/issue-1432.wit

@@ -0,0 +1,19 @@
+//@ async = true
+
+package a:b;
+
+world w {
+    import i;
+    export i;
+}
+
+interface i {
+    use t.{r};
+    f: func(p: stream<r>);
+}
+
+interface t {
+    record r {
+        x: u32,
+    }
+}