Generate nominal interface/type ids in wit-component (#2516)

* Generate nominal interface/type ids in `wit-component`

This commit is a long-overdue change in the componentization process in
`wit-component` to ensure that when a component is generated there are
unique `TypeId` and `InterfaceId` entries corresponding to all
types/interfaces which map to the final component. This is distinct from
the structure of WIT where the same `InterfaceId`, for example, can be
both imported and exported. When generating a component this means that
two distinct interfaces are imported/exported. Previously
`wit-component` has had a lot of very carefully filled out and handled
maps and such to ensure everything works, and this should make things
significantly easier because the possibility over overlap is now
nonexistent.

* Leverage nominal types/interfaces in encoding

This commit builds on the previous commit to simplify internal
structures within the encoding process of a component. Notably there's
no longer any need to have separate maps for imports/exports and instead
everything is located in a single set of maps. Not major changes, but
this is hoped to unlock future changes and additionally pave the way to
simplifying some internals in the future.

Author: alexcrichton

crates/wit-component/src/encoding.rs

@@ -399,17 +399,15 @@ pub struct EncodingState<'a> {
     /// A map of adapter module instances and the index of their instance.
     adapter_instances: IndexMap<&'a str, u32>,
 
-    /// Imported instances and what index they were imported as.
-    imported_instances: IndexMap<InterfaceId, u32>,
+    /// Imported/exported instances and what index they were imported as.
+    instances: IndexMap<InterfaceId, u32>,
     imported_funcs: IndexMap<String, u32>,
-    exported_instances: IndexMap<InterfaceId, u32>,
 
     /// Maps used when translating types to the component model binary format.
     /// Note that imports and exports are stored in separate maps since they
     /// need fresh hierarchies of types in case the same interface is both
     /// imported and exported.
-    import_type_encoding_maps: TypeEncodingMaps<'a>,
-    export_type_encoding_maps: TypeEncodingMaps<'a>,
+    type_encoding_maps: TypeEncodingMaps<'a>,
 
     /// Cache of items that have been aliased from core instances.
     ///
@@ -566,7 +564,7 @@ impl<'a> EncodingState<'a> {
         let instance_idx = self
             .component
             .import(name, ComponentTypeRef::Instance(instance_type_idx));
-        let prev = self.imported_instances.insert(interface_id, instance_idx);
+        let prev = self.instances.insert(interface_id, instance_idx);
         assert!(prev.is_none());
         Ok(())
     }
@@ -600,18 +598,10 @@ impl<'a> EncodingState<'a> {
         Ok(())
     }
 
-    fn alias_imported_type(&mut self, interface: InterfaceId, id: TypeId) -> u32 {
+    fn alias_instance_type_export(&mut self, interface: InterfaceId, id: TypeId) -> u32 {
         let ty = &self.info.encoder.metadata.resolve.types[id];
         let name = ty.name.as_ref().expect("type must have a name");
-        let instance = self.imported_instances[&interface];
-        self.component
-            .alias_export(instance, name, ComponentExportKind::Type)
-    }
-
-    fn alias_exported_type(&mut self, interface: InterfaceId, id: TypeId) -> u32 {
-        let ty = &self.info.encoder.metadata.resolve.types[id];
-        let name = ty.name.as_ref().expect("type must have a name");
-        let instance = self.exported_instances[&interface];
+        let instance = self.instances[&interface];
         self.component
             .alias_export(instance, name, ComponentExportKind::Type)
     }
@@ -674,9 +664,9 @@ impl<'a> EncodingState<'a> {
             // If this resource is owned by a world then it's a top-level
             // resource which means it must have already been translated so
             // it's available for lookup in `import_type_map`.
-            TypeOwner::World(_) => self.import_type_encoding_maps.id_to_index[&id],
+            TypeOwner::World(_) => self.type_encoding_maps.id_to_index[&id],
             TypeOwner::Interface(i) => {
-                let instance = self.imported_instances[&i];
+                let instance = self.instances[&i];
                 let name = ty.name.as_ref().expect("resources must be named");
                 self.component
                     .alias_export(instance, name, ComponentExportKind::Type)
@@ -886,17 +876,7 @@ impl<'a> EncodingState<'a> {
         nested.type_encoding_maps.def_to_index.clear();
         for (name, idx) in nested.imports.drain(..) {
             let id = reverse_map[&idx];
-            let owner = match resolve.types[id].owner {
-                TypeOwner::Interface(id) => id,
-                _ => unreachable!(),
-            };
-            let idx = if owner == export || exports_used.contains(&owner) {
-                log::trace!("consulting exports for {id:?}");
-                nested.state.export_type_encoding_maps.id_to_index[&id]
-            } else {
-                log::trace!("consulting imports for {id:?}");
-                nested.state.import_type_encoding_maps.id_to_index[&id]
-            };
+            let idx = nested.state.type_encoding_maps.id_to_index[&id];
             imports.push((name, ComponentExportKind::Type, idx))
         }
 
@@ -966,7 +946,7 @@ impl<'a> EncodingState<'a> {
             instance_index,
             None,
         );
-        let prev = self.exported_instances.insert(export, idx);
+        let prev = self.instances.insert(export, idx);
         assert!(prev.is_none());
 
         // After everything is all said and done remove all the type information
@@ -977,8 +957,8 @@ impl<'a> EncodingState<'a> {
         // necessary via aliases from the exported instance which is the new
         // source of truth for all these types.
         for (_name, id) in resolve.interfaces[export].types.iter() {
-            self.export_type_encoding_maps.id_to_index.remove(id);
-            self.export_type_encoding_maps
+            self.type_encoding_maps.id_to_index.remove(id);
+            self.type_encoding_maps
                 .def_to_index
                 .remove(&resolve.types[*id].kind);
         }
@@ -1296,7 +1276,7 @@ impl<'a> EncodingState<'a> {
                     let ((name, _), _) = interface.lowerings.get_index(*index).unwrap();
                     let func_index = match &interface.interface {
                         Some(interface_id) => {
-                            let instance_index = self.imported_instances[interface_id];
+                            let instance_index = self.instances[interface_id];
                             self.component.alias_export(
                                 instance_index,
                                 name,
@@ -1602,7 +1582,7 @@ impl<'a> EncodingState<'a> {
                             dtor,
                         );
                         let prev = self
-                            .export_type_encoding_maps
+                            .type_encoding_maps
                             .id_to_index
                             .insert(*ty, resource_idx);
                         assert!(prev.is_none());
@@ -1759,19 +1739,19 @@ impl<'a> EncodingState<'a> {
             Import::ExportedResourceDrop(_key, id) => {
                 let index = self
                     .component
-                    .resource_drop(self.export_type_encoding_maps.id_to_index[id]);
+                    .resource_drop(self.type_encoding_maps.id_to_index[id]);
                 Ok((ExportKind::Func, index))
             }
             Import::ExportedResourceRep(_key, id) => {
                 let index = self
                     .component
-                    .resource_rep(self.export_type_encoding_maps.id_to_index[id]);
+                    .resource_rep(self.type_encoding_maps.id_to_index[id]);
                 Ok((ExportKind::Func, index))
             }
             Import::ExportedResourceNew(_key, id) => {
                 let index = self
                     .component
-                    .resource_new(self.export_type_encoding_maps.id_to_index[id]);
+                    .resource_new(self.type_encoding_maps.id_to_index[id]);
                 Ok((ExportKind::Func, index))
             }
 
@@ -2078,7 +2058,7 @@ impl<'a> EncodingState<'a> {
             Lowering::Direct => {
                 let func_index = match &import.interface {
                     Some(interface) => {
-                        let instance_index = self.imported_instances[interface];
+                        let instance_index = self.instances[interface];
                         self.component
                             .alias_export(instance_index, name, ComponentExportKind::Func)
                     }
@@ -3336,6 +3316,8 @@ impl ComponentEncoder {
                 .merge_world_imports_based_on_semver(self.metadata.world)?;
         }
 
+        self.finalize_resolve_with_nominal_ids();
+
         let world = ComponentWorld::new(self).context("failed to decode world from module")?;
         let mut state = EncodingState {
             component: ComponentBuilder::default(),
@@ -3346,11 +3328,9 @@ impl ComponentEncoder {
             fixups_module_index: None,
             adapter_modules: IndexMap::new(),
             adapter_instances: IndexMap::new(),
-            import_type_encoding_maps: Default::default(),
-            export_type_encoding_maps: Default::default(),
-            imported_instances: Default::default(),
+            type_encoding_maps: Default::default(),
+            instances: Default::default(),
             imported_funcs: Default::default(),
-            exported_instances: Default::default(),
             aliased_core_items: Default::default(),
             info: &world,
             export_task_initialization_wrappers: HashMap::new(),
@@ -3376,6 +3356,63 @@ impl ComponentEncoder {
 
         Ok(bytes)
     }
+
+    /// Call the `generate_nominal_type_ids` method on the `Resolve` that we're
+    /// using, adjusting any preexisting keys/pointers as necessary.
+    ///
+    /// This is the final step after merging all known `Resolve`s together
+    /// before a component is actually created. By creating a unique
+    /// `InterfaceId` for all interfaces it makes the generation process easier
+    /// since there's no need to fret about whether an `InterfaceId` is an
+    /// import or an export for example.
+    fn finalize_resolve_with_nominal_ids(&mut self) {
+        // Before calling `generate_nominal_type_ids` we need to handle the fact
+        // that the exports of the world are going to be rewritten. The only
+        // pointers we have into those are the exports of the main module and
+        // adapters. To handle this, before we generate nominal ids, indices of
+        // exports are saved here on the stack to get restored later on.
+        // Effectively we're clearing out the exports and rebuilding them later.
+        let world = &self.metadata.resolve.worlds[self.metadata.world];
+        let main_module_exports = self
+            .main_module_exports
+            .iter()
+            .map(|i| world.exports.get_index_of(i).unwrap())
+            .collect::<Vec<_>>();
+        let adapter_exports = self
+            .adapters
+            .values()
+            .map(|adapter| {
+                adapter
+                    .required_exports
+                    .iter()
+                    .map(|i| world.exports.get_index_of(i).unwrap())
+                    .collect::<Vec<_>>()
+            })
+            .collect::<Vec<_>>();
+
+        // With everything saved this will modify `Resolve` to ensure there's a
+        // nominal identifier for all interfaces (e.g. not both simultaneously
+        // imported and exported).
+        self.metadata
+            .resolve
+            .generate_nominal_type_ids(self.metadata.world);
+
+        // Rebuild the sets of exports now that the world's exports have been
+        // clobbered.
+        self.main_module_exports.clear();
+        let world = &self.metadata.resolve.worlds[self.metadata.world];
+        for index in main_module_exports {
+            let (key, _) = world.exports.get_index(index).unwrap();
+            self.main_module_exports.insert(key.clone());
+        }
+        for (exports, adapter) in adapter_exports.into_iter().zip(self.adapters.values_mut()) {
+            adapter.required_exports.clear();
+            for index in exports {
+                let (key, _) = world.exports.get_index(index).unwrap();
+                adapter.required_exports.insert(key.clone());
+            }
+        }
+    }
 }
 
 impl ComponentWorld<'_> {

crates/wit-component/src/encoding/types.rs

@@ -474,22 +474,10 @@ impl<'a> ValtypeEncoder<'a> for RootTypeEncoder<'_, 'a> {
             .import(name, ComponentTypeRef::Type(TypeBounds::SubResource))
     }
     fn import_type(&mut self, interface: InterfaceId, id: TypeId) -> u32 {
-        if !self.import_types {
-            if let Some(cur) = self.interface {
-                let set = &self.state.info.exports_used[&cur];
-                if set.contains(&interface) {
-                    return self.state.alias_exported_type(interface, id);
-                }
-            }
-        }
-        self.state.alias_imported_type(interface, id)
+        self.state.alias_instance_type_export(interface, id)
     }
     fn type_encoding_maps(&mut self) -> &mut TypeEncodingMaps<'a> {
-        if self.import_types {
-            &mut self.state.import_type_encoding_maps
-        } else {
-            &mut self.state.export_type_encoding_maps
-        }
+        &mut self.state.type_encoding_maps
     }
 }
 
@@ -528,7 +516,7 @@ impl<'a> ValtypeEncoder<'a> for InstanceTypeEncoder<'_, 'a> {
     fn import_type(&mut self, interface: InterfaceId, id: TypeId) -> u32 {
         self.ty.alias(Alias::Outer {
             count: 1,
-            index: self.state.alias_imported_type(interface, id),
+            index: self.state.alias_instance_type_export(interface, id),
             kind: ComponentOuterAliasKind::Type,
         });
         self.ty.type_count() - 1

crates/wit-component/src/validation.rs

@@ -2397,10 +2397,27 @@ impl NameMangling for Legacy {
                 .interfaces
                 .get(name.interface().as_str())
             {
+                // If the interface from the package is directly in `items` then
+                // return that.
                 let key = WorldKey::Interface(*id);
                 if items.contains_key(&key) {
                     return Ok((key, *id));
                 }
+
+                // .. otherwise see if any interface in `items` is a clone of
+                // the package's interface. This means it's created by
+                // `generate_nominal_type_ids` and is used to match up exports
+                // to their nominal clone since the original is no longer
+                // exported.
+                for k in items.keys() {
+                    let i = match *k {
+                        WorldKey::Interface(id) => id,
+                        WorldKey::Name(_) => continue,
+                    };
+                    if resolve.interfaces[i].clone_of == Some(*id) {
+                        return Ok((WorldKey::Interface(i), i));
+                    }
+                }
             }
         }
 

crates/wit-component/tests/components/async-builtins/component.wat

@@ -239,9 +239,8 @@
   (alias core export $main "[async-lift-stackful]foo" (core func $"[async-lift-stackful]foo" (;33;)))
   (func $foo (;0;) (type 0) (canon lift (core func $"[async-lift-stackful]foo") (memory $memory) (realloc $realloc) string-encoding=utf8 async))
   (export $"#func1 foo" (@name "foo") (;1;) "foo" (func $foo))
-  (type (;1;) (func (param "s" string) (result string)))
   (alias core export $main "[async-lift-stackful]foo:foo/bar#foo" (core func $"[async-lift-stackful]foo:foo/bar#foo" (;34;)))
-  (func $"#func2 foo" (@name "foo") (;2;) (type 1) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#foo") (memory $memory) (realloc $realloc) string-encoding=utf8 async))
+  (func $"#func2 foo" (@name "foo") (;2;) (type 0) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#foo") (memory $memory) (realloc $realloc) string-encoding=utf8 async))
   (component $foo:foo/bar-shim-component (;0;)
     (type (;0;) (func (param "s" string) (result string)))
     (import "import-func-foo" (func (;0;) (type 0)))

crates/wit-component/tests/components/async-export-with-callback/component.wat

@@ -38,10 +38,9 @@
   (alias core export $main "[callback][async-lift]foo" (core func $"[callback][async-lift]foo" (;2;)))
   (func $foo (;0;) (type 0) (canon lift (core func $"[async-lift]foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async (callback $"[callback][async-lift]foo")))
   (export $"#func1 foo" (@name "foo") (;1;) "foo" (func $foo))
-  (type (;1;) (func (param "s" string) (result string)))
   (alias core export $main "[async-lift]foo:foo/bar#foo" (core func $"[async-lift]foo:foo/bar#foo" (;3;)))
   (alias core export $main "[callback][async-lift]foo:foo/bar#foo" (core func $"[callback][async-lift]foo:foo/bar#foo" (;4;)))
-  (func $"#func2 foo" (@name "foo") (;2;) (type 1) (canon lift (core func $"[async-lift]foo:foo/bar#foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async (callback $"[callback][async-lift]foo:foo/bar#foo")))
+  (func $"#func2 foo" (@name "foo") (;2;) (type 0) (canon lift (core func $"[async-lift]foo:foo/bar#foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async (callback $"[callback][async-lift]foo:foo/bar#foo")))
   (component $foo:foo/bar-shim-component (;0;)
     (type (;0;) (func (param "s" string) (result string)))
     (import "import-func-foo" (func (;0;) (type 0)))

crates/wit-component/tests/components/async-export/component.wat

@@ -40,15 +40,14 @@
   (alias core export $main "cabi_realloc" (core func $cabi_realloc (;1;)))
   (func $foo (;0;) (type 3) (canon lift (core func $"[async-lift-stackful]foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async))
   (export $"#func1 foo" (@name "foo") (;1;) "foo" (func $foo))
-  (type (;4;) (func (param "s" string) (result string)))
   (alias core export $main "[async-lift-stackful]foo:foo/bar#foo" (core func $"[async-lift-stackful]foo:foo/bar#foo" (;2;)))
-  (func $"#func2 foo" (@name "foo") (;2;) (type 4) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async))
-  (type (;5;) (func (result string)))
+  (func $"#func2 foo" (@name "foo") (;2;) (type 3) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#foo") (memory $memory) (realloc $cabi_realloc) string-encoding=utf8 async))
+  (type (;4;) (func (result string)))
   (alias core export $main "[async-lift-stackful]foo:foo/bar#get-string" (core func $"[async-lift-stackful]foo:foo/bar#get-string" (;3;)))
-  (func $get-string (;3;) (type 5) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#get-string") (memory $memory) string-encoding=utf8 async))
-  (type (;6;) (func (result 2)))
+  (func $get-string (;3;) (type 4) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#get-string") (memory $memory) string-encoding=utf8 async))
+  (type (;5;) (func (result 2)))
   (alias core export $main "[async-lift-stackful]foo:foo/bar#get-big" (core func $"[async-lift-stackful]foo:foo/bar#get-big" (;4;)))
-  (func $get-big (;4;) (type 6) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#get-big") (memory $memory) async))
+  (func $get-big (;4;) (type 5) (canon lift (core func $"[async-lift-stackful]foo:foo/bar#get-big") (memory $memory) async))
   (component $foo:foo/bar-shim-component (;0;)
     (type (;0;) (func (param "s" string) (result string)))
     (import "import-func-foo" (func (;0;) (type 0)))