Canonicalize types for runtime usage in ExternType::from_wasmtime

This is required to avoid an assertion failure in the regression test for the
previous commit. In order to canonicalize for runtime usage, we need the
associated `TypeCollection`, which required plumbing it through a bunch of
places in the code base.

crates/wasmtime/src/runtime/code.rs

@@ -26,7 +26,7 @@ pub struct CodeObject {
     /// Note that this type has a significant destructor which unregisters
     /// signatures within the `Engine` it was originally tied to, and this ends
     /// up corresponding to the lifetime of a `Component` or `Module`.
-    signatures: TypeCollection,
+    signatures: Arc<TypeCollection>,
 
     /// Type information for the loaded object.
     ///
@@ -36,7 +36,7 @@ pub struct CodeObject {
 }
 
 impl CodeObject {
-    pub fn new(mmap: Arc<CodeMemory>, signatures: TypeCollection, types: Types) -> CodeObject {
+    pub fn new(mmap: Arc<CodeMemory>, signatures: Arc<TypeCollection>, types: Types) -> CodeObject {
         // The corresponding unregister for this is below in `Drop for
         // CodeObject`.
         crate::module::register_code(&mmap);
@@ -61,7 +61,7 @@ impl CodeObject {
         self.types.module_types()
     }
 
-    pub fn signatures(&self) -> &TypeCollection {
+    pub fn signatures(&self) -> &Arc<TypeCollection> {
         &self.signatures
     }
 }

crates/wasmtime/src/runtime/component/component.rs

@@ -368,6 +368,7 @@ impl Component {
         let resources = Arc::new(PrimaryMap::new());
         f(&InstanceType {
             types: self.types(),
+            signatures: self.signatures(),
             resources: &resources,
         })
     }
@@ -407,6 +408,7 @@ impl Component {
         // Assemble the `CodeObject` artifact which is shared by all core wasm
         // modules as well as the final component.
         let types = Arc::new(types);
+        let signatures = Arc::new(signatures);
         let code = Arc::new(CodeObject::new(code_memory, signatures, types.into()));
 
         // Convert all information about static core wasm modules into actual
@@ -453,7 +455,7 @@ impl Component {
         self.inner.component_types()
     }
 
-    pub(crate) fn signatures(&self) -> &TypeCollection {
+    pub(crate) fn signatures(&self) -> &Arc<TypeCollection> {
         self.inner.code.signatures()
     }
 
@@ -798,6 +800,10 @@ impl ComponentRuntimeInfo for ComponentInner {
         }
     }
 
+    fn component_signatures(&self) -> &Arc<TypeCollection> {
+        self.code.signatures()
+    }
+
     fn realloc_func_type(&self) -> &Arc<dyn Any + Send + Sync> {
         &self.realloc_func_type
     }

crates/wasmtime/src/runtime/component/linker.rs

@@ -161,6 +161,7 @@ impl<T> Linker<T> {
     fn typecheck<'a>(&'a self, component: &'a Component) -> Result<TypeChecker<'a>> {
         let mut cx = TypeChecker {
             types: component.types(),
+            signatures: component.signatures(),
             strings: &self.strings,
             imported_resources: Default::default(),
         };
@@ -185,6 +186,7 @@ impl<T> Linker<T> {
             component.ty(),
             &InstanceType {
                 types: cx.types,
+                signatures: component.signatures(),
                 resources: &cx.imported_resources,
             },
         ))

crates/wasmtime/src/runtime/component/matching.rs

@@ -3,6 +3,7 @@ use crate::component::linker::{Definition, Strings};
 use crate::component::ResourceType;
 use crate::prelude::*;
 use crate::runtime::vm::component::ComponentInstance;
+use crate::type_registry::TypeCollection;
 use crate::types::matching;
 use crate::Module;
 use alloc::sync::Arc;
@@ -15,6 +16,7 @@ use wasmtime_environ::PrimaryMap;
 
 pub struct TypeChecker<'a> {
     pub types: &'a Arc<ComponentTypes>,
+    pub signatures: &'a Arc<TypeCollection>,
     pub strings: &'a Strings,
     pub imported_resources: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
 }
@@ -23,6 +25,7 @@ pub struct TypeChecker<'a> {
 #[doc(hidden)]
 pub struct InstanceType<'a> {
     pub types: &'a Arc<ComponentTypes>,
+    pub signatures: &'a Arc<TypeCollection>,
     pub resources: &'a Arc<PrimaryMap<ResourceIndex, ResourceType>>,
 }
 
@@ -167,6 +170,7 @@ impl TypeChecker<'_> {
     fn func(&self, expected: TypeFuncIndex, actual: &HostFunc) -> Result<()> {
         let instance_type = InstanceType {
             types: self.types,
+            signatures: self.signatures,
             resources: &self.imported_resources,
         };
         actual.typecheck(expected, &instance_type)
@@ -188,6 +192,7 @@ impl<'a> InstanceType<'a> {
     pub fn new(instance: &'a ComponentInstance) -> InstanceType<'a> {
         InstanceType {
             types: instance.component_types(),
+            signatures: instance.component_signatures(),
             resources: downcast_arc_ref(instance.resource_types()),
         }
     }

crates/wasmtime/src/runtime/component/types.rs

@@ -1,6 +1,7 @@
 //! This module defines the `Type` type, representing the dynamic form of a component interface type.
 
 use crate::component::matching::InstanceType;
+use crate::type_registry::TypeCollection;
 use crate::{Engine, ExternType, FuncType};
 use alloc::sync::Arc;
 use core::fmt;
@@ -37,6 +38,7 @@ pub use crate::component::resources::ResourceType;
 struct Handle<T> {
     index: T,
     types: Arc<ComponentTypes>,
+    signatures: Arc<TypeCollection>,
     resources: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
 }
 
@@ -45,13 +47,15 @@ impl<T> Handle<T> {
         Handle {
             index,
             types: ty.types.clone(),
+            signatures: ty.signatures.clone(),
             resources: ty.resources.clone(),
         }
     }
 
     fn instance(&self) -> InstanceType<'_> {
         InstanceType {
             types: &self.types,
+            signatures: &self.signatures,
             resources: &self.resources,
         }
     }
@@ -747,7 +751,12 @@ impl Module {
             .map(|((namespace, name), ty)| {
                 (
                     (namespace.as_str(), name.as_str()),
-                    ExternType::from_wasmtime(engine, self.0.types.module_types(), ty),
+                    ExternType::from_wasmtime(
+                        engine,
+                        self.0.types.module_types(),
+                        &self.0.signatures,
+                        ty,
+                    ),
                 )
             })
     }
@@ -760,7 +769,12 @@ impl Module {
         self.0.types[self.0.index].exports.iter().map(|(name, ty)| {
             (
                 name.as_str(),
-                ExternType::from_wasmtime(engine, self.0.types.module_types(), ty),
+                ExternType::from_wasmtime(
+                    engine,
+                    self.0.types.module_types(),
+                    &self.0.signatures,
+                    ty,
+                ),
             )
         })
     }

crates/wasmtime/src/runtime/module.rs

@@ -492,6 +492,7 @@ impl Module {
 
         // Package up all our data into a `CodeObject` and delegate to the final
         // step of module compilation.
+        let signatures = Arc::new(signatures);
         let code = Arc::new(CodeObject::new(code_memory, signatures, types.into()));
         Module::from_parts_raw(engine, code, info, true)
     }
@@ -705,14 +706,15 @@ impl Module {
     ) -> impl ExactSizeIterator<Item = ImportType<'module>> + 'module {
         let module = self.compiled_module().module();
         let types = self.types();
+        let signatures = self.signatures();
         let engine = self.engine();
         module
             .imports()
             .map(move |(imp_mod, imp_field, mut ty)| {
                 ty.canonicalize_for_runtime_usage(&mut |i| {
                     self.signatures().shared_type(i).unwrap()
                 });
-                ImportType::new(imp_mod, imp_field, ty, types, engine)
+                ImportType::new(imp_mod, imp_field, ty, types, signatures, engine)
             })
             .collect::<Vec<_>>()
             .into_iter()
@@ -777,9 +779,16 @@ impl Module {
     ) -> impl ExactSizeIterator<Item = ExportType<'module>> + 'module {
         let module = self.compiled_module().module();
         let types = self.types();
+        let signatures = self.signatures();
         let engine = self.engine();
         module.exports.iter().map(move |(name, entity_index)| {
-            ExportType::new(name, module.type_of(*entity_index), types, engine)
+            ExportType::new(
+                name,
+                module.type_of(*entity_index),
+                types,
+                signatures,
+                engine,
+            )
         })
     }
 
@@ -832,6 +841,7 @@ impl Module {
         Some(ExternType::from_wasmtime(
             self.engine(),
             self.types(),
+            self.signatures(),
             &module.type_of(*entity_index),
         ))
     }

crates/wasmtime/src/runtime/types.rs

@@ -1,4 +1,8 @@
 use crate::prelude::*;
+use crate::{
+    type_registry::{RegisteredType, TypeCollection},
+    Engine,
+};
 use core::fmt::{self, Display, Write};
 use wasmtime_environ::{
     EngineOrModuleTypeIndex, EntityType, Global, IndexType, Limits, Memory, ModuleTypes, Table,
@@ -7,8 +11,6 @@ use wasmtime_environ::{
     WasmSubType, WasmValType,
 };
 
-use crate::{type_registry::RegisteredType, Engine};
-
 pub(crate) mod matching;
 
 // Type Representations
@@ -1194,6 +1196,7 @@ impl ExternType {
     pub(crate) fn from_wasmtime(
         engine: &Engine,
         types: &ModuleTypes,
+        signatures: &TypeCollection,
         ty: &EntityType,
     ) -> ExternType {
         match ty {
@@ -1202,7 +1205,10 @@ impl ExternType {
                     FuncType::from_shared_type_index(engine, *e).into()
                 }
                 EngineOrModuleTypeIndex::Module(m) => {
-                    let subty = &types[*m];
+                    let mut subty = types[*m].clone();
+                    subty.canonicalize_for_runtime_usage(&mut |idx| {
+                        signatures.shared_type(idx).unwrap()
+                    });
                     FuncType::from_wasm_func_type(
                         engine,
                         subty.is_final,
@@ -2906,6 +2912,7 @@ pub struct ImportType<'module> {
     /// The type of the import.
     ty: EntityType,
     types: &'module ModuleTypes,
+    signatures: &'module TypeCollection,
     engine: &'module Engine,
 }
 
@@ -2917,6 +2924,7 @@ impl<'module> ImportType<'module> {
         name: &'module str,
         ty: EntityType,
         types: &'module ModuleTypes,
+        signatures: &'module TypeCollection,
         engine: &'module Engine,
     ) -> ImportType<'module> {
         assert!(ty.is_canonicalized_for_runtime_usage());
@@ -2925,6 +2933,7 @@ impl<'module> ImportType<'module> {
             name,
             ty,
             types,
+            signatures,
             engine,
         }
     }
@@ -2942,7 +2951,7 @@ impl<'module> ImportType<'module> {
 
     /// Returns the expected type of this import.
     pub fn ty(&self) -> ExternType {
-        ExternType::from_wasmtime(self.engine, self.types, &self.ty)
+        ExternType::from_wasmtime(self.engine, self.types, self.signatures, &self.ty)
     }
 }
 
@@ -2972,6 +2981,7 @@ pub struct ExportType<'module> {
     /// The type of the export.
     ty: EntityType,
     types: &'module ModuleTypes,
+    signatures: &'module TypeCollection,
     engine: &'module Engine,
 }
 
@@ -2982,12 +2992,14 @@ impl<'module> ExportType<'module> {
         name: &'module str,
         ty: EntityType,
         types: &'module ModuleTypes,
+        signatures: &'module TypeCollection,
         engine: &'module Engine,
     ) -> ExportType<'module> {
         ExportType {
             name,
             ty,
             types,
+            signatures,
             engine,
         }
     }
@@ -2999,7 +3011,7 @@ impl<'module> ExportType<'module> {
 
     /// Returns the type of this export.
     pub fn ty(&self) -> ExternType {
-        ExternType::from_wasmtime(self.engine, self.types, &self.ty)
+        ExternType::from_wasmtime(self.engine, self.types, self.signatures, &self.ty)
     }
 }
 

crates/wasmtime/src/runtime/vm/component.rs

@@ -11,6 +11,7 @@ use crate::runtime::vm::{
     SendSyncPtr, VMArrayCallFunction, VMFuncRef, VMGlobalDefinition, VMMemoryDefinition,
     VMOpaqueContext, VMStore, VMStoreRawPtr, VMWasmCallFunction, ValRaw, VmPtr, VmSafe,
 };
+use crate::type_registry::TypeCollection;
 use alloc::alloc::Layout;
 use alloc::sync::Arc;
 use core::any::Any;
@@ -533,6 +534,11 @@ impl ComponentInstance {
         self.runtime_info.component_types()
     }
 
+    /// Returns the type collection that this instance is associated with.
+    pub fn component_signatures(&self) -> &Arc<TypeCollection> {
+        self.runtime_info.component_signatures()
+    }
+
     /// Get the canonical ABI's `realloc` function's runtime type.
     pub fn realloc_func_ty(&self) -> &Arc<dyn Any + Send + Sync> {
         self.runtime_info.realloc_func_type()
@@ -908,6 +914,9 @@ pub trait ComponentRuntimeInfo: Send + Sync + 'static {
     /// Returns a handle to the tables of type information for this component.
     fn component_types(&self) -> &Arc<ComponentTypes>;
 
+    /// Returns a handle to the type collection for this component.
+    fn component_signatures(&self) -> &Arc<TypeCollection>;
+
     /// Get the `wasmtime::FuncType` for the canonical ABI's `realloc` function.
     fn realloc_func_type(&self) -> &Arc<dyn Any + Send + Sync>;
 }