Implement the module linking alias section (#2451)

This commit is intended to do almost everything necessary for processing
the alias section of module linking. Most of this is internal
refactoring, the highlights being:

* Type contents are now stored separately from a `wasmtime_env::Module`.
  Given that modules can freely alias types and have them used all over
  the place, it seemed best to have one canonical location to type
  storage which everywhere else points to (with indices). A new
  `TypeTables` structure is produced during compilation which is shared
  amongst all member modules in a wasm blob.

* Instantiation is heavily refactored to account for module linking. The
  main gotcha here is that imports are now listed as "initializers". We
  have a sort of pseudo-bytecode-interpreter which interprets the
  initialization of a module. This is more complicated than just
  matching imports at this point because in the module linking proposal
  the module, alias, import, and instance sections may all be
  interleaved. This means that imports aren't guaranteed to show up at
  the beginning of the address space for modules/instances.

Otherwise most of the changes here largely fell out from these two
design points. Aliases are recorded as initializers in this scheme.
Copying around type information and/or just knowing type information
during compilation is also pretty easy since everything is just a
pointer into a `TypeTables` and we don't have to actually copy any types
themselves. Lots of various refactorings were necessary to accomodate
these changes.

Tests are hoped to cover a breadth of functionality here, but not
necessarily a depth. There's still one more piece of the module linking
proposal missing which is exporting instances/modules, which will come
in a future PR.

It's also worth nothing that there's one large TODO which isn't
implemented in this change that I plan on opening an issue for.
With module linking when a set of modules comes back from compilation
each modules has all the trampolines for the entire set of modules. This
is quite a lot of duplicate trampolines across module-linking modules.
We'll want to refactor this at some point to instead have only one set
of trampolines per set of module linking modules and have them shared
from there. I figured it was best to separate out this change, however,
since it's purely related to resource usage, and doesn't impact
non-module-linking modules at all.

cc #2094

cranelift/wasm/src/environ/mod.rs

@@ -6,6 +6,6 @@ mod spec;
 
 pub use crate::environ::dummy::DummyEnvironment;
 pub use crate::environ::spec::{
-    FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode, TargetEnvironment, WasmError,
-    WasmFuncType, WasmResult, WasmType,
+    Alias, FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode, TargetEnvironment,
+    WasmError, WasmFuncType, WasmResult, WasmType,
 };

cranelift/wasm/src/environ/spec.rs

@@ -9,7 +9,8 @@
 use crate::state::FuncTranslationState;
 use crate::translation_utils::{
     DataIndex, ElemIndex, EntityIndex, EntityType, Event, EventIndex, FuncIndex, Global,
-    GlobalIndex, Memory, MemoryIndex, ModuleIndex, Table, TableIndex, TypeIndex,
+    GlobalIndex, InstanceIndex, InstanceTypeIndex, Memory, MemoryIndex, ModuleIndex,
+    ModuleTypeIndex, SignatureIndex, Table, TableIndex, TypeIndex,
 };
 use core::convert::From;
 use core::convert::TryFrom;
@@ -202,6 +203,30 @@ pub enum ReturnMode {
     FallthroughReturn,
 }
 
+/// An entry in the alias section of a wasm module (from the module linking
+/// proposal)
+pub enum Alias {
+    /// A parent's module is being aliased into our own index space.
+    ///
+    /// Note that the index here is in the parent's index space, not our own.
+    ParentModule(ModuleIndex),
+
+    /// A parent's type is being aliased into our own index space
+    ///
+    /// Note that the index here is in the parent's index space, not our own.
+    ParentType(TypeIndex),
+
+    /// A previously created instance is having one of its exports aliased into
+    /// our index space.
+    Child {
+        /// The index we're aliasing.
+        instance: InstanceIndex,
+        /// The nth export that we're inserting into our own index space
+        /// locally.
+        export: usize,
+    },
+}
+
 /// Environment affecting the translation of a WebAssembly.
 pub trait TargetEnvironment {
     /// Get the information needed to produce Cranelift IR for the given target.
@@ -684,6 +709,27 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         Err(WasmError::Unsupported("module linking".to_string()))
     }
 
+    /// Translates a type index to its signature index, only called for type
+    /// indices which point to functions.
+    fn type_to_signature(&self, index: TypeIndex) -> WasmResult<SignatureIndex> {
+        drop(index);
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
+    /// Translates a type index to its module type index, only called for type
+    /// indices which point to modules.
+    fn type_to_module_type(&self, index: TypeIndex) -> WasmResult<ModuleTypeIndex> {
+        drop(index);
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
+    /// Translates a type index to its instance type index, only called for type
+    /// indices which point to instances.
+    fn type_to_instance_type(&self, index: TypeIndex) -> WasmResult<InstanceTypeIndex> {
+        drop(index);
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
     /// Provides the number of imports up front. By default this does nothing, but
     /// implementations can use this to preallocate memory if desired.
     fn reserve_imports(&mut self, _num: u32) -> WasmResult<()> {
@@ -845,6 +891,22 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         name: &'data str,
     ) -> WasmResult<()>;
 
+    /// Declares an instance export to the environment.
+    fn declare_instance_export(
+        &mut self,
+        index: InstanceIndex,
+        name: &'data str,
+    ) -> WasmResult<()> {
+        drop((index, name));
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
+    /// Declares an instance export to the environment.
+    fn declare_module_export(&mut self, index: ModuleIndex, name: &'data str) -> WasmResult<()> {
+        drop((index, name));
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
     /// Notifies the implementation that all exports have been declared.
     fn finish_exports(&mut self) -> WasmResult<()> {
         Ok(())
@@ -952,6 +1014,12 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         drop(amount);
     }
 
+    /// Declares that a module will come later with the type signature provided.
+    fn declare_module(&mut self, ty: TypeIndex) -> WasmResult<()> {
+        drop(ty);
+        Err(WasmError::Unsupported("module linking".to_string()))
+    }
+
     /// Called at the beginning of translating a module.
     ///
     /// The `index` argument is a monotonically increasing index which
@@ -982,4 +1050,14 @@ pub trait ModuleEnvironment<'data>: TargetEnvironment {
         drop((module, args));
         Err(WasmError::Unsupported("wasm instance".to_string()))
     }
+
+    /// Declares a new alias being added to this module.
+    ///
+    /// The alias comes from the `instance` specified (or the parent if `None`
+    /// is supplied) and the index is either in the module's own index spaces
+    /// for the parent or an index into the exports for nested instances.
+    fn declare_alias(&mut self, alias: Alias) -> WasmResult<()> {
+        drop(alias);
+        Err(WasmError::Unsupported("wasm alias".to_string()))
+    }
 }

cranelift/wasm/src/lib.rs

@@ -57,7 +57,7 @@ mod state;
 mod translation_utils;
 
 pub use crate::environ::{
-    DummyEnvironment, FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode,
+    Alias, DummyEnvironment, FuncEnvironment, GlobalVariable, ModuleEnvironment, ReturnMode,
     TargetEnvironment, WasmError, WasmFuncType, WasmResult, WasmType,
 };
 pub use crate::func_translator::FuncTranslator;

cranelift/wasm/src/module_translator.rs

@@ -2,10 +2,10 @@
 //! to deal with each part of it.
 use crate::environ::{ModuleEnvironment, WasmResult};
 use crate::sections_translator::{
-    parse_data_section, parse_element_section, parse_event_section, parse_export_section,
-    parse_function_section, parse_global_section, parse_import_section, parse_instance_section,
-    parse_memory_section, parse_name_section, parse_start_section, parse_table_section,
-    parse_type_section,
+    parse_alias_section, parse_data_section, parse_element_section, parse_event_section,
+    parse_export_section, parse_function_section, parse_global_section, parse_import_section,
+    parse_instance_section, parse_memory_section, parse_module_section, parse_name_section,
+    parse_start_section, parse_table_section, parse_type_section,
 };
 use crate::state::ModuleTranslationState;
 use cranelift_codegen::timing;
@@ -113,15 +113,15 @@ pub fn translate_module<'data>(
 
             Payload::ModuleSection(s) => {
                 validator.module_section(&s)?;
-                environ.reserve_modules(s.get_count());
+                parse_module_section(s, environ)?;
             }
             Payload::InstanceSection(s) => {
                 validator.instance_section(&s)?;
                 parse_instance_section(s, environ)?;
             }
             Payload::AliasSection(s) => {
                 validator.alias_section(&s)?;
-                unimplemented!("module linking not implemented yet")
+                parse_alias_section(s, environ)?;
             }
             Payload::ModuleCodeSectionStart {
                 count,

cranelift/wasm/src/sections_translator.rs

@@ -7,7 +7,7 @@
 //! The special case of the initialize expressions for table elements offsets or global variables
 //! is handled, according to the semantics of WebAssembly, to only specific expressions that are
 //! interpreted on the fly.
-use crate::environ::{ModuleEnvironment, WasmError, WasmResult};
+use crate::environ::{Alias, ModuleEnvironment, WasmError, WasmResult};
 use crate::state::ModuleTranslationState;
 use crate::translation_utils::{
     tabletype_to_type, type_to_type, DataIndex, ElemIndex, EntityIndex, EntityType, Event,
@@ -36,9 +36,15 @@ fn entity_type(
     environ: &mut dyn ModuleEnvironment<'_>,
 ) -> WasmResult<EntityType> {
     Ok(match ty {
-        ImportSectionEntryType::Function(sig) => EntityType::Function(TypeIndex::from_u32(sig)),
-        ImportSectionEntryType::Module(sig) => EntityType::Module(TypeIndex::from_u32(sig)),
-        ImportSectionEntryType::Instance(sig) => EntityType::Instance(TypeIndex::from_u32(sig)),
+        ImportSectionEntryType::Function(sig) => {
+            EntityType::Function(environ.type_to_signature(TypeIndex::from_u32(sig))?)
+        }
+        ImportSectionEntryType::Module(sig) => {
+            EntityType::Module(environ.type_to_module_type(TypeIndex::from_u32(sig))?)
+        }
+        ImportSectionEntryType::Instance(sig) => {
+            EntityType::Instance(environ.type_to_instance_type(TypeIndex::from_u32(sig))?)
+        }
         ImportSectionEntryType::Memory(ty) => EntityType::Memory(memory(ty)),
         ImportSectionEntryType::Event(evt) => EntityType::Event(event(evt)),
         ImportSectionEntryType::Global(ty) => {
@@ -156,24 +162,40 @@ pub fn parse_import_section<'data>(
 
     for entry in imports {
         let import = entry?;
-        match entity_type(import.ty, environ)? {
-            EntityType::Function(idx) => {
-                environ.declare_func_import(idx, import.module, import.field)?;
+        match import.ty {
+            ImportSectionEntryType::Function(sig) => {
+                environ.declare_func_import(
+                    TypeIndex::from_u32(sig),
+                    import.module,
+                    import.field,
+                )?;
             }
-            EntityType::Module(idx) => {
-                environ.declare_module_import(idx, import.module, import.field)?;
+            ImportSectionEntryType::Module(sig) => {
+                environ.declare_module_import(
+                    TypeIndex::from_u32(sig),
+                    import.module,
+                    import.field,
+                )?;
+            }
+            ImportSectionEntryType::Instance(sig) => {
+                environ.declare_instance_import(
+                    TypeIndex::from_u32(sig),
+                    import.module,
+                    import.field,
+                )?;
             }
-            EntityType::Instance(idx) => {
-                environ.declare_instance_import(idx, import.module, import.field)?;
+            ImportSectionEntryType::Memory(ty) => {
+                environ.declare_memory_import(memory(ty), import.module, import.field)?;
             }
-            EntityType::Memory(ty) => {
-                environ.declare_memory_import(ty, import.module, import.field)?;
+            ImportSectionEntryType::Event(e) => {
+                environ.declare_event_import(event(e), import.module, import.field)?;
             }
-            EntityType::Event(e) => environ.declare_event_import(e, import.module, import.field)?,
-            EntityType::Global(ty) => {
+            ImportSectionEntryType::Global(ty) => {
+                let ty = global(ty, environ, GlobalInit::Import)?;
                 environ.declare_global_import(ty, import.module, import.field)?;
             }
-            EntityType::Table(ty) => {
+            ImportSectionEntryType::Table(ty) => {
+                let ty = table(ty, environ)?;
                 environ.declare_table_import(ty, import.module, import.field)?;
             }
         }
@@ -316,9 +338,15 @@ pub fn parse_export_section<'data>(
             ExternalKind::Global => {
                 environ.declare_global_export(GlobalIndex::new(index), field)?
             }
-            ExternalKind::Type | ExternalKind::Module | ExternalKind::Instance => {
-                unimplemented!("module linking not implemented yet")
+            ExternalKind::Module => {
+                environ.declare_module_export(ModuleIndex::new(index), field)?
             }
+            ExternalKind::Instance => {
+                environ.declare_instance_export(InstanceIndex::new(index), field)?
+            }
+
+            // this never gets past validation
+            ExternalKind::Type => unreachable!(),
         }
     }
 
@@ -476,12 +504,25 @@ pub fn parse_name_section<'data>(
     Ok(())
 }
 
+/// Parses the Module section of the wasm module.
+pub fn parse_module_section<'data>(
+    section: wasmparser::ModuleSectionReader<'data>,
+    environ: &mut dyn ModuleEnvironment<'data>,
+) -> WasmResult<()> {
+    environ.reserve_modules(section.get_count());
+
+    for module_ty in section {
+        environ.declare_module(TypeIndex::from_u32(module_ty?))?;
+    }
+    Ok(())
+}
+
 /// Parses the Instance section of the wasm module.
 pub fn parse_instance_section<'data>(
     section: wasmparser::InstanceSectionReader<'data>,
     environ: &mut dyn ModuleEnvironment<'data>,
 ) -> WasmResult<()> {
-    environ.reserve_types(section.get_count())?;
+    environ.reserve_instances(section.get_count());
 
     for instance in section {
         let instance = instance?;
@@ -509,3 +550,29 @@ pub fn parse_instance_section<'data>(
     }
     Ok(())
 }
+
+/// Parses the Alias section of the wasm module.
+pub fn parse_alias_section<'data>(
+    section: wasmparser::AliasSectionReader<'data>,
+    environ: &mut dyn ModuleEnvironment<'data>,
+) -> WasmResult<()> {
+    for alias in section {
+        let alias = alias?;
+        let alias = match alias.instance {
+            wasmparser::AliasedInstance::Parent => {
+                match alias.kind {
+                    ExternalKind::Module => Alias::ParentModule(ModuleIndex::from_u32(alias.index)),
+                    ExternalKind::Type => Alias::ParentType(TypeIndex::from_u32(alias.index)),
+                    // shouldn't get past validation
+                    _ => unreachable!(),
+                }
+            }
+            wasmparser::AliasedInstance::Child(i) => Alias::Child {
+                instance: InstanceIndex::from_u32(i),
+                export: alias.index as usize,
+            },
+        };
+        environ.declare_alias(alias)?;
+    }
+    Ok(())
+}

cranelift/wasm/src/translation_utils.rs

@@ -97,6 +97,18 @@ entity_impl!(InstanceIndex);
 pub struct EventIndex(u32);
 entity_impl!(EventIndex);
 
+/// Specialized index for just module types.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct ModuleTypeIndex(u32);
+entity_impl!(ModuleTypeIndex);
+
+/// Specialized index for just instance types.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
+#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
+pub struct InstanceTypeIndex(u32);
+entity_impl!(InstanceTypeIndex);
+
 /// An index of an entity.
 #[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
@@ -131,13 +143,13 @@ pub enum EntityType {
     Table(Table),
     /// A function type where the index points to the type section and records a
     /// function signature.
-    Function(TypeIndex),
+    Function(SignatureIndex),
     /// An instance where the index points to the type section and records a
     /// instance's exports.
-    Instance(TypeIndex),
+    Instance(InstanceTypeIndex),
     /// A module where the index points to the type section and records a
     /// module's imports and exports.
-    Module(TypeIndex),
+    Module(ModuleTypeIndex),
 }
 
 /// A WebAssembly global.

crates/cranelift/src/func_environ.rs

@@ -1039,7 +1039,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         callee: ir::Value,
         call_args: &[ir::Value],
     ) -> WasmResult<ir::Inst> {
-        let sig_index = self.module.types[ty_index].unwrap_function();
         let pointer_type = self.pointer_type();
 
         let table_entry_addr = pos.ins().table_addr(pointer_type, table, callee, 0);
@@ -1071,7 +1070,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
                 let vmctx = self.vmctx(pos.func);
                 let base = pos.ins().global_value(pointer_type, vmctx);
                 let offset =
-                    i32::try_from(self.offsets.vmctx_vmshared_signature_id(sig_index)).unwrap();
+                    i32::try_from(self.offsets.vmctx_vmshared_signature_id(ty_index)).unwrap();
 
                 // Load the caller ID.
                 let mut mem_flags = ir::MemFlags::trusted();