add CloneMaps type and parameter to Resolve::merge_worlds (#2361)

This allows the caller to determine which types and/or interfaces were cloned as
part of the merge operation.

Note that I made this an `&mut` parameter to `Resolve::merge_worlds` rather than
a return value in order to make it a bit more efficient when calling
`Resolve::merge_worlds` multiple times and collecting the clones across all
those calls.

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

Author: dicej

crates/wit-component/src/metadata.rs

@@ -50,7 +50,7 @@ use wasm_encoder::{
 };
 use wasm_metadata::Producers;
 use wasmparser::{BinaryReader, Encoding, Parser, Payload};
-use wit_parser::{Package, PackageName, Resolve, World, WorldId, WorldItem, WorldKey};
+use wit_parser::{CloneMaps, Package, PackageName, Resolve, World, WorldId, WorldItem, WorldKey};
 
 const CURRENT_VERSION: u8 = 0x04;
 const CUSTOM_SECTION_NAME: &str = "wit-component-encoding";
@@ -413,7 +413,7 @@ impl Bindgen {
         let world = remap.map_world(world, None)?;
         let exports = self.resolve.worlds[world].exports.keys().cloned().collect();
         self.resolve
-            .merge_worlds(world, self.world)
+            .merge_worlds(world, self.world, &mut CloneMaps::default())
             .context("failed to merge worlds from two documents")?;
 
         self.metadata.import_encodings.merge(import_encodings)?;

crates/wit-parser/src/resolve.rs

@@ -23,6 +23,8 @@ use crate::{
     UnresolvedPackageGroup, World, WorldId, WorldItem, WorldKey, WorldSpan,
 };
 
+pub use clone::CloneMaps;
+
 mod clone;
 
 /// Representation of a fully resolved set of WIT packages.
@@ -866,8 +868,17 @@ package {name} is defined in two different locations:\n\
     /// producing a singular world that will be the final component's
     /// interface.
     ///
+    /// During the merge operation, some of the types and/or interfaces in
+    /// `from` might need to be cloned so that backreferences point to `into`
+    /// instead of `from`.  Any such clones will be added to `clone_maps`.
+    ///
     /// This operation can fail if the imports/exports overlap.
-    pub fn merge_worlds(&mut self, from: WorldId, into: WorldId) -> Result<()> {
+    pub fn merge_worlds(
+        &mut self,
+        from: WorldId,
+        into: WorldId,
+        clone_maps: &mut CloneMaps,
+    ) -> Result<()> {
         let mut new_imports = Vec::new();
         let mut new_exports = Vec::new();
 
@@ -960,9 +971,11 @@ package {name} is defined in two different locations:\n\
         let mut cloner = clone::Cloner::new(self, TypeOwner::World(from), TypeOwner::World(into));
         cloner.register_world_type_overlap(from, into);
         for (name, item) in new_imports.iter_mut().chain(&mut new_exports) {
-            cloner.world_item(name, item);
+            cloner.world_item(name, item, clone_maps);
         }
 
+        clone_maps.types.extend(cloner.types);
+
         // Insert any new imports and new exports found first.
         let into_world = &mut self.worlds[into];
         for (name, import) in new_imports {
@@ -3654,7 +3667,7 @@ impl Remap {
                 // in the function itself.
                 let mut new_item = item.1.clone();
                 let key = WorldKey::Name(n.clone());
-                cloner.world_item(&key, &mut new_item);
+                cloner.world_item(&key, &mut new_item, &mut CloneMaps::default());
                 match &mut new_item {
                     WorldItem::Function(f) => f.name = n.clone(),
                     WorldItem::Type(id) => cloner.resolve.types[*id].name = Some(n.clone()),

crates/wit-parser/src/resolve/clone.rs

@@ -20,6 +20,30 @@
 use crate::*;
 use std::collections::HashMap;
 
+/// Represents the results of cloning types and/or interfaces as part of a
+/// `Resolve::merge_worlds` operation.
+#[derive(Default)]
+pub struct CloneMaps {
+    pub(super) types: HashMap<TypeId, TypeId>,
+    pub(super) interfaces: HashMap<InterfaceId, InterfaceId>,
+}
+
+impl CloneMaps {
+    /// The types cloned during a `Resolve::merge_worlds` operation.
+    ///
+    /// The key is the original type, and the value is the clone.
+    pub fn types(&self) -> &HashMap<TypeId, TypeId> {
+        &self.types
+    }
+
+    /// The interfaces cloned during a `Resolve::merge_worlds` operation.
+    ///
+    /// The key is the original interface, and the value is the clone.
+    pub fn interfaces(&self) -> &HashMap<InterfaceId, InterfaceId> {
+        &self.interfaces
+    }
+}
+
 pub struct Cloner<'a> {
     pub resolve: &'a mut Resolve,
     prev_owner: TypeOwner,
@@ -28,7 +52,7 @@ pub struct Cloner<'a> {
     /// This map keeps track, in the current scope of types, of all copied
     /// types. This deduplicates copying types to ensure that they're only
     /// copied at most once.
-    types: HashMap<TypeId, TypeId>,
+    pub types: HashMap<TypeId, TypeId>,
 
     /// If `None` then it's inferred from `self.new_owner`.
     pub new_package: Option<PackageId>,
@@ -64,7 +88,7 @@ impl<'a> Cloner<'a> {
         }
     }
 
-    pub fn world_item(&mut self, key: &WorldKey, item: &mut WorldItem) {
+    pub fn world_item(&mut self, key: &WorldKey, item: &mut WorldItem, clone_maps: &mut CloneMaps) {
         match key {
             WorldKey::Name(_) => {}
             WorldKey::Interface(_) => return,
@@ -78,7 +102,9 @@ impl<'a> Cloner<'a> {
                 self.function(f);
             }
             WorldItem::Interface { id, .. } => {
-                self.interface(id);
+                let old = *id;
+                self.interface(id, &mut clone_maps.types);
+                clone_maps.interfaces.insert(old, *id);
             }
         }
     }
@@ -171,7 +197,7 @@ impl<'a> Cloner<'a> {
         }
     }
 
-    fn interface(&mut self, id: &mut InterfaceId) {
+    fn interface(&mut self, id: &mut InterfaceId, cloned_types: &mut HashMap<TypeId, TypeId>) {
         let mut new = self.resolve.interfaces[*id].clone();
         let next_id = self.resolve.interfaces.next_id();
         let mut clone = Cloner::new(
@@ -185,6 +211,7 @@ impl<'a> Cloner<'a> {
         for func in new.functions.values_mut() {
             clone.function(func);
         }
+        cloned_types.extend(clone.types);
         new.package = Some(self.new_package.unwrap_or_else(|| match self.new_owner {
             TypeOwner::Interface(id) => self.resolve.interfaces[id].package.unwrap(),
             TypeOwner::World(id) => self.resolve.worlds[id].package.unwrap(),