Repository: bytecodealliance/wasm-tools
Pull request #1792 (2024-09-17 20:37:28)

 Maintain more internal consistency checks when merging worlds 
This commit fixes two issues, both preexisting, which represented
violated invariants of a `Resolve` after worlds were merged. World
merging is used during component generation and might be used for
bindings generation so it's important to maintain the various dynamic
invariants of `Resolve` to ensure that surrounding tooling works. An
example of this is that printing a merged world doesn't produce
parseable WIT today because some of the internal invariants aren't
respected.

This commit beefs up the `Resolve::assert_valid` function and then
applies necessary fixes to resolve these assertion failures on
preexisting test cases. The main difference is that anonymous
interfaces/types are now "cloned" when they move from one world to
another. This is required to ensure that all the various links between
these items are consistent and point to the right place. This cloning
infrastructure may end up being more generalized in the future if the
need arises, but for now it's quite limited.

alexcrichton: I'll note that this is stacked on https://github.com/bytecodealliance/wasm-tools/pull/1791 so the first commit can be ignored.

Patch:
diff --git a/crates/wit-parser/src/resolve.rs b/crates/wit-parser/src/resolve.rs
index 3a9ee24e..bd03376f 100644
--- a/crates/wit-parser/src/resolve.rs
+++ b/crates/wit-parser/src/resolve.rs
@@ -19,6 +19,8 @@ use crate::{
     WorldKey, WorldSpan,
 };
 
+mod clone;
+
 /// Representation of a fully resolved set of WIT packages.
 ///
 /// This structure contains a graph of WIT packages and all of their contents
@@ -747,6 +749,8 @@ package {name} is defined in two different locations:\n\
         let from_world = &self.worlds[from];
         let into_world = &self.worlds[into];
 
+        log::trace!("merging {} into {}", from_world.name, into_world.name);
+
         // First walk over all the imports of `from` world and figure out what
         // to do with them.
         //
@@ -755,13 +759,15 @@ package {name} is defined in two different locations:\n\
         // same. Otherwise queue up a new import since if `from` has more
         // imports than `into` then it's fine to add new imports.
         for (name, from_import) in from_world.imports.iter() {
+            let name_str = self.name_world_key(name);
             match into_world.imports.get(name) {
                 Some(into_import) => {
-                    let name = self.name_world_key(name);
+                    log::trace!("info/from shared import on `{name_str}`");
                     self.merge_world_item(from_import, into_import)
-                        .with_context(|| format!("failed to merge world import {name}"))?;
+                        .with_context(|| format!("failed to merge world import {name_str}"))?;
                 }
                 None => {
+                    log::trace!("new import: `{name_str}`");
                     new_imports.push((name.clone(), from_import.clone()));
                 }
             }
@@ -788,13 +794,15 @@ package {name} is defined in two different locations:\n\
         // Next walk over exports of `from` and process these similarly to
         // imports.
         for (name, from_export) in from_world.exports.iter() {
+            let name_str = self.name_world_key(name);
             match into_world.exports.get(name) {
                 Some(into_export) => {
-                    let name = self.name_world_key(name);
+                    log::trace!("info/from shared export on `{name_str}`");
                     self.merge_world_item(from_export, into_export)
-                        .with_context(|| format!("failed to merge world export {name}"))?;
+                        .with_context(|| format!("failed to merge world export {name_str}"))?;
                 }
                 None => {
+                    log::trace!("new export `{name_str}`");
                     // See comments in `ensure_can_add_world_export` for why
                     // this is slightly different than imports.
                     self.ensure_can_add_world_export(
@@ -811,14 +819,33 @@ package {name} is defined in two different locations:\n\
             }
         }
 
+        // For all the new imports and exports they may need to be "cloned" to
+        // be able to belong to the new world. For example:
+        //
+        // * Anonymous interfaces have a `package` field which points to the
+        //   package of the containing world, but `from` and `into` may not be
+        //   in the same package.
+        //
+        // * Type imports have an `owner` field that point to `from`, but they
+        //   now need to point to `into` instead.
+        //
+        // Cloning is no trivial task, however, so cloning is delegated to a
+        // submodule to perform a "deep" clone and copy items into new arena
+        // entries as necessary.
+        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);
+        }
+
         // Insert any new imports and new exports found first.
-        let into = &mut self.worlds[into];
+        let into_world = &mut self.worlds[into];
         for (name, import) in new_imports {
-            let prev = into.imports.insert(name, import);
+            let prev = into_world.imports.insert(name, import);
             assert!(prev.is_none());
         }
         for (name, export) in new_exports {
-            let prev = into.exports.insert(name, export);
+            let prev = into_world.exports.insert(name, export);
             assert!(prev.is_none());
         }
 
@@ -1428,7 +1455,13 @@ package {name} is defined in two different locations:\n\
             for (name, item) in world.imports.iter().chain(world.exports.iter()) {
                 log::debug!("validating world item: {}", self.name_world_key(name));
                 match item {
-                    WorldItem::Interface { .. } => {}
+                    WorldItem::Interface { id, .. } => {
+                        // anonymous interfaces must belong to the same package
+                        // as the world's package.
+                        if matches!(name, WorldKey::Name(_)) {
+                            assert_eq!(self.interfaces[*id].package, world.package);
+                        }
+                    }
                     WorldItem::Function(f) => {
                         assert!(!matches!(name, WorldKey::Interface(_)));
                         assert_eq!(f.name, name.clone().unwrap_name());
@@ -1438,12 +1471,7 @@ package {name} is defined in two different locations:\n\
                         assert!(types.insert(*ty));
                         let ty = &self.types[*ty];
                         assert_eq!(ty.name, Some(name.clone().unwrap_name()));
-
-                        // TODO: `Resolve::merge_worlds` doesn't uphold this
-                        // invariant, and that should be fixed.
-                        if false {
-                            assert_eq!(ty.owner, TypeOwner::World(id));
-                        }
+                        assert_eq!(ty.owner, TypeOwner::World(id));
                     }
                 }
             }
diff --git a/crates/wit-parser/src/resolve/clone.rs b/crates/wit-parser/src/resolve/clone.rs
new file mode 100644
index 00000000..21885df5
--- /dev/null
+++ b/crates/wit-parser/src/resolve/clone.rs
@@ -0,0 +1,199 @@
+//! A helper, non public, module to assist with cloning items within a
+//! `Resolve`.
+//!
+//! Cloning items is not as simple as calling `Clone` due to the nature of how
+//! `TypeId` tracks relationships between interfaces and types. A full deep
+//! clone requires walking the full structure and allocating new id links. This
+//! is akin to, for example, creating a deep copy of an `Rc<T>` by calling
+//! `Clone for T`.
+//!
+//! This is currently used when merging worlds together to help copy anonymously
+//! named items from one world to another.
+//!
+//! The general structure of this module is that each method takes a mutable
+//! reference to an AST item and updates it as necessary internally, delegating
+//! to other methods for internal AST items.
+//!
+//! This module does not at the time of this writing have full support for
+//! cloning everything within a `Resolve`.
+
+use crate::*;
+use std::collections::HashMap;
+
+pub struct Cloner<'a> {
+    resolve: &'a mut Resolve,
+    prev_owner: TypeOwner,
+    new_owner: TypeOwner,
+
+    /// 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>,
+}
+
+impl<'a> Cloner<'a> {
+    pub fn new(
+        resolve: &'a mut Resolve,
+        prev_owner: TypeOwner,
+        new_owner: TypeOwner,
+    ) -> Cloner<'a> {
+        Cloner {
+            prev_owner,
+            new_owner,
+            resolve,
+            types: Default::default(),
+        }
+    }
+
+    pub fn register_world_type_overlap(&mut self, from: WorldId, into: WorldId) {
+        let into = &self.resolve.worlds[into];
+        let from = &self.resolve.worlds[from];
+        for (name, into_import) in into.imports.iter() {
+            let WorldKey::Name(_) = name else { continue };
+            let WorldItem::Type(into_id) = into_import else {
+                continue;
+            };
+            let Some(WorldItem::Type(from_id)) = from.imports.get(name) else {
+                continue;
+            };
+            self.types.insert(*from_id, *into_id);
+        }
+    }
+
+    pub fn world_item(&mut self, key: &WorldKey, item: &mut WorldItem) {
+        match key {
+            WorldKey::Name(_) => {}
+            WorldKey::Interface(_) => return,
+        }
+
+        match item {
+            WorldItem::Type(t) => {
+                self.type_id(t);
+            }
+            WorldItem::Function(f) => {
+                self.function(f);
+            }
+            WorldItem::Interface { id, .. } => {
+                self.interface(id);
+            }
+        }
+    }
+
+    fn type_id(&mut self, ty: &mut TypeId) {
+        if !self.types.contains_key(ty) {
+            let mut new = self.resolve.types[*ty].clone();
+            self.type_def(&mut new);
+            let id = self.resolve.types.alloc(new);
+            self.types.insert(*ty, id);
+        }
+        *ty = self.types[ty];
+    }
+
+    fn type_def(&mut self, def: &mut TypeDef) {
+        if def.owner != TypeOwner::None {
+            assert_eq!(def.owner, self.prev_owner);
+            def.owner = self.new_owner;
+        }
+        match &mut def.kind {
+            TypeDefKind::Type(Type::Id(id)) => {
+                if self.resolve.types[*id].owner == self.prev_owner {
+                    self.type_id(id);
+                } else {
+                    // ..
+                }
+            }
+            TypeDefKind::Type(_)
+            | TypeDefKind::Resource
+            | TypeDefKind::Flags(_)
+            | TypeDefKind::Enum(_) => {}
+            TypeDefKind::Handle(Handle::Own(ty) | Handle::Borrow(ty)) => {
+                self.type_id(ty);
+            }
+            TypeDefKind::Option(ty) | TypeDefKind::List(ty) => {
+                self.ty(ty);
+            }
+            TypeDefKind::Tuple(list) => {
+                for ty in list.types.iter_mut() {
+                    self.ty(ty);
+                }
+            }
+            TypeDefKind::Record(r) => {
+                for field in r.fields.iter_mut() {
+                    self.ty(&mut field.ty);
+                }
+            }
+            TypeDefKind::Variant(r) => {
+                for case in r.cases.iter_mut() {
+                    if let Some(ty) = &mut case.ty {
+                        self.ty(ty);
+                    }
+                }
+            }
+            TypeDefKind::Result(r) => {
+                if let Some(ok) = &mut r.ok {
+                    self.ty(ok);
+                }
+                if let Some(err) = &mut r.err {
+                    self.ty(err);
+                }
+            }
+            TypeDefKind::Future(f) => {
+                if let Some(ty) = f {
+                    self.ty(ty);
+                }
+            }
+            TypeDefKind::Stream(_) => unimplemented!(),
+            TypeDefKind::Unknown => {}
+        }
+    }
+
+    fn ty(&mut self, ty: &mut Type) {
+        match ty {
+            Type::Id(id) => self.type_id(id),
+            _ => {}
+        }
+    }
+
+    fn function(&mut self, func: &mut Function) {
+        match &mut func.kind {
+            FunctionKind::Freestanding => {}
+            FunctionKind::Method(id) | FunctionKind::Static(id) | FunctionKind::Constructor(id) => {
+                self.type_id(id)
+            }
+        }
+        for (_, ty) in func.params.iter_mut() {
+            self.ty(ty);
+        }
+        match &mut func.results {
+            Results::Named(named) => {
+                for (_, ty) in named {
+                    self.ty(ty);
+                }
+            }
+            Results::Anon(ty) => self.ty(ty),
+        }
+    }
+
+    fn interface(&mut self, id: &mut InterfaceId) {
+        let mut new = self.resolve.interfaces[*id].clone();
+        let next_id = self.resolve.interfaces.next_id();
+        let mut clone = Cloner::new(
+            self.resolve,
+            TypeOwner::Interface(*id),
+            TypeOwner::Interface(next_id),
+        );
+        for id in new.types.values_mut() {
+            clone.type_id(id);
+        }
+        for func in new.functions.values_mut() {
+            clone.function(func);
+        }
+        new.package = Some(match self.new_owner {
+            TypeOwner::Interface(id) => self.resolve.interfaces[id].package.unwrap(),
+            TypeOwner::World(id) => self.resolve.worlds[id].package.unwrap(),
+            TypeOwner::None => unreachable!(),
+        });
+        *id = self.resolve.interfaces.alloc(new);
+        assert_eq!(*id, next_id);
+    }
+}


Test patch:
diff --git a/tests/cli/print-interfaces-in-world-correctly.wit b/tests/cli/print-interfaces-in-world-correctly.wit
new file mode 100644
index 00000000..cc090e91
--- /dev/null
+++ b/tests/cli/print-interfaces-in-world-correctly.wit
@@ -0,0 +1,24 @@
+// RUN: component wit %
+
+package foo:bar;
+
+world foo {
+  import a: interface {
+    use local.{t1};
+    use a:b/foreign.{t2};
+  }
+  export a: interface {
+    use local.{t1};
+    use a:b/foreign.{t2};
+  }
+}
+
+interface local {
+  type t1 = u32;
+}
+
+package a:b {
+  interface foreign {
+    type t2 = u32;
+  }
+}
diff --git a/tests/cli/print-interfaces-in-world-correctly.wit.stdout b/tests/cli/print-interfaces-in-world-correctly.wit.stdout
new file mode 100644
index 00000000..56ae5efa
--- /dev/null
+++ b/tests/cli/print-interfaces-in-world-correctly.wit.stdout
@@ -0,0 +1,25 @@
+/// RUN: component wit %
+package foo:bar;
+
+interface local {
+  type t1 = u32;
+}
+
+world foo {
+  import local;
+  import a:b/foreign;
+  import a: interface {
+    use local.{t1};
+    use a:b/foreign.{t2};
+  }
+
+  export a: interface {
+    use local.{t1};
+    use a:b/foreign.{t2};
+  }
+}
+package a:b {
+  interface foreign {
+    type t2 = u32;
+  }
+}
diff --git a/tests/cli/reparent-anonymous-interfaces-on-merge.wit b/tests/cli/reparent-anonymous-interfaces-on-merge.wit
new file mode 100644
index 00000000..62979034
--- /dev/null
+++ b/tests/cli/reparent-anonymous-interfaces-on-merge.wit
@@ -0,0 +1,19 @@
+// RUN: component embed --dummy --world w1 % | \
+//      component embed         --world w2 % | \
+//      component embed         --world w3 % | \
+//      component embed         --world w4 % | \
+//      component wit
+
+package a:b;
+
+world w1 {}
+world w2 {
+  import a: interface {}
+}
+world w3 {
+  export a: interface {}
+}
+world w4 {
+  import b: interface {}
+  export b: interface {}
+}
diff --git a/tests/cli/reparent-anonymous-interfaces-on-merge.wit.stdout b/tests/cli/reparent-anonymous-interfaces-on-merge.wit.stdout
new file mode 100644
index 00000000..2b8ca640
--- /dev/null
+++ b/tests/cli/reparent-anonymous-interfaces-on-merge.wit.stdout
@@ -0,0 +1,32 @@
+package root:root;
+
+world root {
+  import a: interface {
+  }
+  import b: interface {
+  }
+
+  export a: interface {
+  }
+  export b: interface {
+  }
+}
+package a:b {
+  world w1 {
+  }
+  world w2 {
+    import a: interface {
+    }
+  }
+  world w3 {
+    export a: interface {
+    }
+  }
+  world w4 {
+    import b: interface {
+    }
+
+    export b: interface {
+    }
+  }
+}
diff --git a/tests/cli/world-merge-add-resource-func.wit b/tests/cli/world-merge-add-resource-func.wit
new file mode 100644
index 00000000..f1e0e3ef
--- /dev/null
+++ b/tests/cli/world-merge-add-resource-func.wit
@@ -0,0 +1,16 @@
+// RUN: component embed --dummy --world a % | \
+//      component embed         --world b % | \
+//      component wit
+
+
+package a:b;
+
+world a {
+  resource a;
+}
+
+world b {
+  resource a {
+    constructor();
+  }
+}
diff --git a/tests/cli/world-merge-add-resource-func.wit.stdout b/tests/cli/world-merge-add-resource-func.wit.stdout
new file mode 100644
index 00000000..43a3bb6d
--- /dev/null
+++ b/tests/cli/world-merge-add-resource-func.wit.stdout
@@ -0,0 +1,17 @@
+package root:root;
+
+world root {
+  resource a {
+    constructor();
+  }
+}
+package a:b {
+  world a {
+    resource a;
+  }
+  world b {
+    resource a {
+      constructor();
+    }
+  }
+}