Rust: merge structurally equal types in bindgen (#1468)

* merge equal types

* fix

* add test

* Refactor type equality algorithm structure

* Handle `alias == primitive` by juggling typedefs/`Type::Id` a bit
  more carefully.
* Avoid using `_ => false`
* Fix handling of `own`/`borrow` and `future`/`stream` to be more
  uniform like the rest of the algorithm.
* Avoid special-casing resources, but for now consider them always
  not-equal.

* Adjust CLI flag handling

* Refactor how structurally-equal types are generated

Leverage the recent support for nominal type IDs to remove some
now-no-longer-necessary infrastructure. This additionally overrides the
`define_type` method in the Rust generator to handle aliases at one
location instead of in multiple locations.

---------

Co-authored-by: Alex Crichton <alex@alexcrichton.com>

Author: chenyan2002

crates/core/src/lib.rs

@@ -176,25 +176,32 @@ pub trait InterfaceGenerator<'a> {
     }
 
     fn define_type(&mut self, name: &str, id: TypeId) {
-        let ty = &self.resolve().types[id];
-        match &ty.kind {
-            TypeDefKind::Record(record) => self.type_record(id, name, record, &ty.docs),
-            TypeDefKind::Resource => self.type_resource(id, name, &ty.docs),
-            TypeDefKind::Flags(flags) => self.type_flags(id, name, flags, &ty.docs),
-            TypeDefKind::Tuple(tuple) => self.type_tuple(id, name, tuple, &ty.docs),
-            TypeDefKind::Enum(enum_) => self.type_enum(id, name, enum_, &ty.docs),
-            TypeDefKind::Variant(variant) => self.type_variant(id, name, variant, &ty.docs),
-            TypeDefKind::Option(t) => self.type_option(id, name, t, &ty.docs),
-            TypeDefKind::Result(r) => self.type_result(id, name, r, &ty.docs),
-            TypeDefKind::List(t) => self.type_list(id, name, t, &ty.docs),
-            TypeDefKind::Type(t) => self.type_alias(id, name, t, &ty.docs),
-            TypeDefKind::Future(t) => self.type_future(id, name, t, &ty.docs),
-            TypeDefKind::Stream(t) => self.type_stream(id, name, t, &ty.docs),
-            TypeDefKind::Handle(_) => panic!("handle types do not require definition"),
-            TypeDefKind::FixedLengthList(..) => todo!(),
-            TypeDefKind::Map(..) => todo!(),
-            TypeDefKind::Unknown => unreachable!(),
-        }
+        define_type(self, name, id)
+    }
+}
+
+pub fn define_type<'a, T>(generator: &mut T, name: &str, id: TypeId)
+where
+    T: InterfaceGenerator<'a> + ?Sized,
+{
+    let ty = &generator.resolve().types[id];
+    match &ty.kind {
+        TypeDefKind::Record(record) => generator.type_record(id, name, record, &ty.docs),
+        TypeDefKind::Resource => generator.type_resource(id, name, &ty.docs),
+        TypeDefKind::Flags(flags) => generator.type_flags(id, name, flags, &ty.docs),
+        TypeDefKind::Tuple(tuple) => generator.type_tuple(id, name, tuple, &ty.docs),
+        TypeDefKind::Enum(enum_) => generator.type_enum(id, name, enum_, &ty.docs),
+        TypeDefKind::Variant(variant) => generator.type_variant(id, name, variant, &ty.docs),
+        TypeDefKind::Option(t) => generator.type_option(id, name, t, &ty.docs),
+        TypeDefKind::Result(r) => generator.type_result(id, name, r, &ty.docs),
+        TypeDefKind::List(t) => generator.type_list(id, name, t, &ty.docs),
+        TypeDefKind::Type(t) => generator.type_alias(id, name, t, &ty.docs),
+        TypeDefKind::Future(t) => generator.type_future(id, name, t, &ty.docs),
+        TypeDefKind::Stream(t) => generator.type_stream(id, name, t, &ty.docs),
+        TypeDefKind::Handle(_) => panic!("handle types do not require definition"),
+        TypeDefKind::FixedLengthList(..) => todo!(),
+        TypeDefKind::Map(..) => todo!(),
+        TypeDefKind::Unknown => unreachable!(),
     }
 }
 

crates/core/src/types.rs

@@ -5,6 +5,7 @@ use wit_parser::*;
 #[derive(Default)]
 pub struct Types {
     type_info: HashMap<TypeId, TypeInfo>,
+    equal_types: UnionFind,
 }
 
 #[derive(Default, Clone, Copy, Debug)]
@@ -93,6 +94,22 @@ impl Types {
             }
         }
     }
+    pub fn collect_equal_types(&mut self, resolve: &Resolve) {
+        for (i, (ty, _)) in resolve.types.iter().enumerate() {
+            // 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) {
+                    continue;
+                }
+                // The correctness of is_structurally_equal relies on the fact
+                // that resolve.types.iter() is in topological order.
+                if self.is_structurally_equal(resolve, ty, earlier) {
+                    self.equal_types.union(ty, earlier);
+                    break;
+                }
+            }
+        }
+    }
 
     fn type_info_func(&mut self, resolve: &Resolve, func: &Function, import: bool) {
         let mut live = LiveTypes::default();
@@ -233,4 +250,185 @@ impl Types {
             None => TypeInfo::default(),
         }
     }
+
+    fn is_structurally_equal(&mut self, resolve: &Resolve, a: TypeId, b: TypeId) -> bool {
+        let a_def = &resolve.types[a].kind;
+        let b_def = &resolve.types[b].kind;
+        if self.equal_types.find(a) == self.equal_types.find(b) {
+            return true;
+        }
+        match (a_def, b_def) {
+            // Peel off typedef layers and continue recursing.
+            (TypeDefKind::Type(a), _) => self.type_id_equal_to_type(resolve, b, a),
+            (_, TypeDefKind::Type(b)) => self.type_id_equal_to_type(resolve, a, b),
+
+            (TypeDefKind::Record(ra), TypeDefKind::Record(rb)) => {
+                ra.fields.len() == rb.fields.len()
+                  // Fields are ordered in WIT, so record {a: T, b: U} is different from {b: U, a: T}
+                  && ra.fields.iter().zip(rb.fields.iter()).all(|(fa, fb)| {
+                      fa.name == fb.name && self.types_equal(resolve, &fa.ty, &fb.ty)
+                  })
+            }
+            (TypeDefKind::Record(_), _) => false,
+            (TypeDefKind::Variant(va), TypeDefKind::Variant(vb)) => {
+                va.cases.len() == vb.cases.len()
+                    && va.cases.iter().zip(vb.cases.iter()).all(|(ca, cb)| {
+                        ca.name == cb.name && self.optional_types_equal(resolve, &ca.ty, &cb.ty)
+                    })
+            }
+            (TypeDefKind::Variant(_), _) => false,
+            (TypeDefKind::Enum(ea), TypeDefKind::Enum(eb)) => {
+                ea.cases.len() == eb.cases.len()
+                    && ea
+                        .cases
+                        .iter()
+                        .zip(eb.cases.iter())
+                        .all(|(ca, cb)| ca.name == cb.name)
+            }
+            (TypeDefKind::Enum(_), _) => false,
+            (TypeDefKind::Flags(fa), TypeDefKind::Flags(fb)) => {
+                fa.flags.len() == fb.flags.len()
+                    && fa
+                        .flags
+                        .iter()
+                        .zip(fb.flags.iter())
+                        .all(|(fa, fb)| fa.name == fb.name)
+            }
+            (TypeDefKind::Flags(_), _) => false,
+            (TypeDefKind::Tuple(ta), TypeDefKind::Tuple(tb)) => {
+                ta.types.len() == tb.types.len()
+                    && ta
+                        .types
+                        .iter()
+                        .zip(tb.types.iter())
+                        .all(|(a, b)| self.types_equal(resolve, a, b))
+            }
+            (TypeDefKind::Tuple(_), _) => false,
+            (TypeDefKind::List(la), TypeDefKind::List(lb)) => self.types_equal(resolve, la, lb),
+            (TypeDefKind::List(_), _) => false,
+            (TypeDefKind::FixedLengthList(ta, sa), TypeDefKind::FixedLengthList(tb, sb)) => {
+                sa == sb && self.types_equal(resolve, ta, tb)
+            }
+            (TypeDefKind::FixedLengthList(..), _) => false,
+            (TypeDefKind::Option(oa), TypeDefKind::Option(ob)) => self.types_equal(resolve, oa, ob),
+            (TypeDefKind::Option(_), _) => false,
+            (TypeDefKind::Result(ra), TypeDefKind::Result(rb)) => {
+                self.optional_types_equal(resolve, &ra.ok, &rb.ok)
+                    && self.optional_types_equal(resolve, &ra.err, &rb.err)
+            }
+            (TypeDefKind::Result(_), _) => false,
+            (TypeDefKind::Map(ak, av), TypeDefKind::Map(bk, bv)) => {
+                self.types_equal(resolve, ak, bk) && self.types_equal(resolve, av, bv)
+            }
+            (TypeDefKind::Map(..), _) => false,
+            (TypeDefKind::Future(a), TypeDefKind::Future(b)) => {
+                self.optional_types_equal(resolve, a, b)
+            }
+            (TypeDefKind::Future(..), _) => false,
+            (TypeDefKind::Stream(a), TypeDefKind::Stream(b)) => {
+                self.optional_types_equal(resolve, a, b)
+            }
+            (TypeDefKind::Stream(..), _) => false,
+            (TypeDefKind::Handle(a), TypeDefKind::Handle(b)) => match (a, b) {
+                (Handle::Own(a), Handle::Own(b)) | (Handle::Borrow(a), Handle::Borrow(b)) => {
+                    self.is_structurally_equal(resolve, *a, *b)
+                }
+                (Handle::Own(_) | Handle::Borrow(_), _) => false,
+            },
+            (TypeDefKind::Handle(_), _) => false,
+            (TypeDefKind::Unknown, _) => unreachable!(),
+
+            // TODO: for now consider all resources not-equal to each other.
+            // This is because the same type id can be used for both an imported
+            // and exported resource where those should be distinct types.
+            (TypeDefKind::Resource, _) => false,
+        }
+    }
+
+    fn types_equal(&mut self, resolve: &Resolve, a: &Type, b: &Type) -> bool {
+        match (a, b) {
+            // Peel off typedef layers and continue recursing.
+            (Type::Id(a), b) => self.type_id_equal_to_type(resolve, *a, b),
+            (a, Type::Id(b)) => self.type_id_equal_to_type(resolve, *b, a),
+
+            // When both a and b are primitives, they're only equal of
+            // the primitives are the same.
+            (
+                Type::Bool
+                | Type::U8
+                | Type::S8
+                | Type::U16
+                | Type::S16
+                | Type::U32
+                | Type::S32
+                | Type::U64
+                | Type::S64
+                | Type::F32
+                | Type::F64
+                | Type::Char
+                | Type::String
+                | Type::ErrorContext,
+                _,
+            ) => a == b,
+        }
+    }
+
+    fn type_id_equal_to_type(&mut self, resolve: &Resolve, a: TypeId, b: &Type) -> bool {
+        let ak = &resolve.types[a].kind;
+        match (ak, b) {
+            (TypeDefKind::Type(a), b) => self.types_equal(resolve, a, b),
+            (_, Type::Id(b)) => self.is_structurally_equal(resolve, a, *b),
+
+            // Type `a` isn't a typedef, and type `b` is a primitive, so it's no
+            // longer possible for them to be equal.
+            _ => false,
+        }
+    }
+
+    fn optional_types_equal(
+        &mut self,
+        resolve: &Resolve,
+        a: &Option<Type>,
+        b: &Option<Type>,
+    ) -> bool {
+        match (a, b) {
+            (Some(a), Some(b)) => self.types_equal(resolve, a, b),
+            (Some(_), None) | (None, Some(_)) => false,
+            (None, None) => true,
+        }
+    }
+
+    pub fn get_representative_type(&mut self, id: TypeId) -> TypeId {
+        self.equal_types.find(id)
+    }
+}
+
+#[derive(Default)]
+pub struct UnionFind {
+    parent: HashMap<TypeId, TypeId>,
+}
+impl UnionFind {
+    fn find(&mut self, id: TypeId) -> TypeId {
+        // Path compression
+        let parent = self.parent.get(&id).copied().unwrap_or(id);
+        if parent != id {
+            let root = self.find(parent);
+            self.parent.insert(id, root);
+            root
+        } else {
+            id
+        }
+    }
+    fn union(&mut self, a: TypeId, b: TypeId) {
+        let ra = self.find(a);
+        let rb = self.find(b);
+        if ra != rb {
+            // Use smaller id as root for determinism
+            if ra < rb {
+                self.parent.insert(rb, ra);
+            } else {
+                self.parent.insert(ra, rb);
+            }
+        }
+    }
 }

crates/rust/src/interface.rs

@@ -2577,6 +2577,16 @@ impl<'a> wit_bindgen_core::InterfaceGenerator<'a> for InterfaceGenerator<'a> {
         self.print_typedef_record(id, record, docs);
     }
 
+    fn define_type(&mut self, name: &str, id: TypeId) {
+        let equal = self.r#gen.types.get_representative_type(id);
+        if equal == id {
+            wit_bindgen_core::define_type(self, name, id)
+        } else {
+            let docs = &self.resolve.types[id].docs;
+            self.print_typedef_alias(id, &Type::Id(equal), &docs);
+        }
+    }
+
     fn type_resource(&mut self, _id: TypeId, name: &str, docs: &Docs) {
         self.rustdoc(docs);
         let camel = to_upper_camel_case(name);

crates/rust/src/lib.rs

@@ -274,6 +274,18 @@ pub struct Opts {
     #[cfg_attr(feature = "clap", clap(flatten))]
     #[cfg_attr(feature = "serde", serde(flatten))]
     pub async_: AsyncFilterSet,
+
+    /// Find all structurally equal types and only generate one type definition for
+    /// each equivalence class. Other types in the same class will be type aliases to the
+    /// generated type. This avoids clone when converting between types that are
+    /// structurally equal, which is useful when import and export the same interface.
+    ///
+    /// Types containing resource, future, or stream are never considered equal.
+    #[cfg_attr(
+        feature = "clap",
+        arg(long, require_equals = true, value_name = "true|false")
+    )]
+    pub merge_structurally_equal_types: Option<Option<bool>>,
 }
 
 impl Opts {
@@ -283,6 +295,18 @@ impl Opts {
         r.opts = self;
         r
     }
+
+    fn merge_structurally_equal_types(&self) -> bool {
+        const DEFAULT: bool = false;
+        match self.merge_structurally_equal_types {
+            // no option passed, use the default
+            None => DEFAULT,
+            // --merge-structurally-equal-types
+            Some(None) => true,
+            // --merge-structurally-equal-types=val
+            Some(Some(val)) => val,
+        }
+    }
 }
 
 impl RustWasm {
@@ -1082,6 +1106,9 @@ impl WorldGenerator for RustWasm {
                 "//   * disable-run-ctors-once-workaround"
             );
         }
+        if self.opts.merge_structurally_equal_types() {
+            uwriteln!(self.src_preamble, "//   * merge_structurally_equal_types");
+        }
         if let Some(s) = &self.opts.export_macro_name {
             uwriteln!(self.src_preamble, "//   * export-macro-name: {s}");
         }
@@ -1101,6 +1128,9 @@ 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.world = Some(world);
 
         let world = &resolve.worlds[world];
@@ -1206,9 +1236,9 @@ impl WorldGenerator for RustWasm {
         _files: &mut Files,
     ) -> Result<()> {
         let mut to_define = Vec::new();
-        for (name, ty_id) in resolve.interfaces[id].types.iter() {
+        for (ty_name, ty_id) in resolve.interfaces[id].types.iter() {
             let full_name = full_wit_type_name(resolve, *ty_id);
-            to_define.push((name, ty_id));
+            to_define.push((ty_name, ty_id));
             self.generated_types.insert(full_name);
         }
 
@@ -1224,8 +1254,8 @@ impl WorldGenerator for RustWasm {
             return Ok(());
         }
 
-        for (name, ty_id) in to_define {
-            r#gen.define_type(&name, *ty_id);
+        for (ty_name, ty_id) in to_define {
+            r#gen.define_type(&ty_name, *ty_id);
         }
 
         let macro_name =
@@ -1427,7 +1457,11 @@ impl WorldGenerator for RustWasm {
     }
 }
 
-fn compute_module_path(name: &WorldKey, resolve: &Resolve, is_export: bool) -> Vec<String> {
+pub(crate) fn compute_module_path(
+    name: &WorldKey,
+    resolve: &Resolve,
+    is_export: bool,
+) -> Vec<String> {
     let mut path = Vec::new();
     if is_export {
         path.push("exports".to_string());

tests/runtime/rust/equal-types/compose.wac

@@ -0,0 +1,6 @@
+package example:composition;
+
+let host = new test:host { ... };
+let proxy = new test:proxy { ...host, ... };
+let runner = new test:runner { ...proxy, ... };
+export runner...;