module-linking: Implement outer module aliases (#2590)

This commit fully implements outer aliases of the module linking
proposal. Outer aliases can now handle multiple-level-up aliases and now
properly also handle closed-over-values of modules that are either
imported or defined.

The structure of `wasmtime::Module` was altered as part of this commit.
It is now a compiled module plus two lists of "upvars", or closed over
values used when instantiating the module. One list of upvars is
compiled artifacts which are submodules that could be used. Another is
module values that are injected via outer aliases. Serialization and
such have been updated as appropriate to handle this.

crates/environ/src/module.rs

@@ -239,11 +239,31 @@ pub enum Initializer {
         args: IndexMap<String, EntityIndex>,
     },
 
-    /// A module is defined into the module index space, and which module is
-    /// being defined is specified by the index payload.
+    /// A module is being created from a set of compiled artifacts.
+    CreateModule {
+        /// The index of the artifact that's being convereted into a module.
+        artifact_index: usize,
+        /// The list of artifacts that this module value will be inheriting.
+        artifacts: Vec<usize>,
+        /// The list of modules that this module value will inherit.
+        modules: Vec<ModuleUpvar>,
+    },
+
+    /// A module is created from a closed-over-module value, defined when this
+    /// module was created.
     DefineModule(usize),
 }
 
+/// Where module values can come from when creating a new module from a compiled
+/// artifact.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ModuleUpvar {
+    /// A module value is inherited from the module creating the new module.
+    Inherit(usize),
+    /// A module value comes from the instance-to-be-created module index space.
+    Local(ModuleIndex),
+}
+
 impl Module {
     /// Allocates the module data structures.
     pub fn new() -> Self {

crates/environ/src/module_environ.rs

@@ -1,6 +1,6 @@
 use crate::module::{
-    Initializer, InstanceSignature, MemoryPlan, Module, ModuleSignature, ModuleType, TableElements,
-    TablePlan, TypeTables,
+    Initializer, InstanceSignature, MemoryPlan, Module, ModuleSignature, ModuleType, ModuleUpvar,
+    TableElements, TablePlan, TypeTables,
 };
 use crate::tunables::Tunables;
 use cranelift_codegen::ir;
@@ -75,6 +75,14 @@ pub struct ModuleTranslation<'data> {
     code_index: u32,
 
     implicit_instances: HashMap<&'data str, InstanceIndex>,
+
+    /// The artifacts which are needed from the parent module when this module
+    /// is created. This is used to insert into `Initializer::CreateModule` when
+    /// this module is defined in the parent.
+    creation_artifacts: Vec<usize>,
+
+    /// Same as `creation_artifacts`, but for modules instead of artifacts.
+    creation_modules: Vec<ModuleUpvar>,
 }
 
 /// Contains function data: byte code and its offset in the module.
@@ -326,8 +334,7 @@ impl<'data> ModuleEnvironment<'data> {
         }
     }
 
-    fn gen_type_of_module(&mut self, module: usize) -> ModuleTypeIndex {
-        let module = &self.results[module].module;
+    fn gen_type_of_module(&mut self, module: &Module) -> ModuleTypeIndex {
         let imports = module
             .imports()
             .map(|(s, field, ty)| {
@@ -880,20 +887,43 @@ and for re-adding support for interface types you can see this issue:
     }
 
     fn module_end(&mut self) {
-        let (record_initializer, done) = match self.in_progress.pop() {
+        self.result.creation_artifacts.shrink_to_fit();
+        self.result.creation_modules.shrink_to_fit();
+
+        let (record_initializer, mut done) = match self.in_progress.pop() {
             Some(m) => (true, mem::replace(&mut self.result, m)),
             None => (false, mem::take(&mut self.result)),
         };
-        self.results.push(done);
+
         if record_initializer {
-            let index = self.results.len() - 1;
+            // Record the type of the module we just finished in our own
+            // module's list of modules.
+            let sig = self.gen_type_of_module(&done.module);
+            self.result.module.modules.push(sig);
+
+            // The root module will store the artifacts for this finished
+            // module at `artifact_index`. This then needs to be inherited by
+            // all later modules coming down to our now-current `self.result`...
+            let mut artifact_index = self.results.len();
+            for result in self.in_progress.iter_mut().chain(Some(&mut self.result)) {
+                result.creation_artifacts.push(artifact_index);
+                artifact_index = result.creation_artifacts.len() - 1;
+            }
+            // ... and then `self.result` needs to create a new module with
+            // whatever was record to save off as its own artifacts/modules.
             self.result
                 .module
                 .initializers
-                .push(Initializer::DefineModule(index));
-            let sig = self.gen_type_of_module(index);
-            self.result.module.modules.push(sig);
+                .push(Initializer::CreateModule {
+                    artifact_index,
+                    artifacts: mem::take(&mut done.creation_artifacts),
+                    modules: mem::take(&mut done.creation_modules),
+                });
         }
+
+        // And the final step is to insert the module into the list of finished
+        // modules to get returned at the end.
+        self.results.push(done);
     }
 
     fn reserve_instances(&mut self, amt: u32) {
@@ -936,16 +966,44 @@ and for re-adding support for interface types you can see this issue:
                 self.result.module.types.push(ty);
             }
 
-            // FIXME(WebAssembly/module-linking#28) unsure how to implement this
-            // at this time, if we can alias imported modules it's a lot harder,
-            // otherwise we'll need to figure out how to translate `index` to a
-            // `usize` for a defined module (creating Initializer::DefineModule)
+            // Modules are a bit trickier since we need to record how to track
+            // the state from the original module down to our own.
             Alias::OuterModule {
                 relative_depth,
                 index,
             } => {
-                drop((relative_depth, index));
-                unimplemented!()
+                // First we can copy the type from the parent module into our
+                // own module to record what type our module definition will
+                // have.
+                let module_idx = self.in_progress.len() - 1 - (relative_depth as usize);
+                let module_ty = self.in_progress[module_idx].module.modules[index];
+                self.result.module.modules.push(module_ty);
+
+                // Next we'll be injecting a module value that is closed over,
+                // and that will be used to define the module into the index
+                // space. Record an initializer about where our module is
+                // sourced from (which will be stored within each module value
+                // itself).
+                let module_index = self.result.creation_modules.len();
+                self.result
+                    .module
+                    .initializers
+                    .push(Initializer::DefineModule(module_index));
+
+                // And finally we need to record a breadcrumb trail of how to
+                // get the module value into `module_index`. The module just
+                // after our destination module will use a `ModuleIndex` to
+                // fetch the module value, and everything else inbetween will
+                // inherit that module's closed-over value.
+                let mut upvar = ModuleUpvar::Local(index);
+                for outer in self.in_progress[module_idx + 1..].iter_mut() {
+                    let upvar = mem::replace(
+                        &mut upvar,
+                        ModuleUpvar::Inherit(outer.creation_modules.len()),
+                    );
+                    outer.creation_modules.push(upvar);
+                }
+                self.result.creation_modules.push(upvar);
             }
 
             // This case is slightly more involved, we'll be recording all the

crates/jit/src/instantiate.rs

@@ -221,7 +221,7 @@ impl CompiledModule {
         artifacts: Vec<CompilationArtifacts>,
         isa: &dyn TargetIsa,
         profiler: &dyn ProfilingAgent,
-    ) -> Result<Vec<Self>, SetupError> {
+    ) -> Result<Vec<Arc<Self>>, SetupError> {
         maybe_parallel!(artifacts.(into_iter | into_par_iter))
             .map(|a| CompiledModule::from_artifacts(a, isa, profiler))
             .collect()
@@ -232,7 +232,7 @@ impl CompiledModule {
         artifacts: CompilationArtifacts,
         isa: &dyn TargetIsa,
         profiler: &dyn ProfilingAgent,
-    ) -> Result<Self, SetupError> {
+    ) -> Result<Arc<Self>, SetupError> {
         // Allocate all of the compiled functions into executable memory,
         // copying over their contents.
         let (code_memory, code_range, finished_functions, trampolines) = build_code_memory(
@@ -266,7 +266,7 @@ impl CompiledModule {
 
         let finished_functions = FinishedFunctions(finished_functions);
 
-        Ok(Self {
+        Ok(Arc::new(Self {
             module: Arc::new(artifacts.module.clone()),
             artifacts,
             code: Arc::new(ModuleCode {
@@ -275,7 +275,7 @@ impl CompiledModule {
             }),
             finished_functions,
             trampolines,
-        })
+        }))
     }
 
     /// Crate an `Instance` from this `CompiledModule`.

crates/wasmtime/src/instance.rs

@@ -62,46 +62,27 @@ fn instantiate(
                     .with_context(|| format!("incompatible import type for `{}`", name))?;
             }
 
-            // Turns out defining any kind of module is pretty easy, we're just
-            // slinging around pointers.
-            Initializer::DefineModule(idx) => {
-                imports.modules.push(module.submodule(*idx));
-            }
-
             // Here we lookup our instance handle, find the right export,
             // and then push that item into our own index space. We eschew
-            // type-checking since only valid modules reach this point.
-            //
-            // Note that export lookup here needs to happen by name. The
-            // `export` index is an index into our local type definition of the
-            // type of the instance to figure out what name it was assigned.
-            // This is where the subtyping happens!
-            //
-            // Note that the unsafety here is because we're asserting that the
-            // handle comes from our same store, but this should be true because
-            // we acquired the handle from an instance in the store.
+            // type-checking since only valid modules should reach this point.
             Initializer::AliasInstanceExport { instance, export } => {
                 let export = &imports.instances[*instance][export];
                 let item = unsafe { Extern::from_wasmtime_export(export, store) };
                 imports.push_extern(&item);
             }
 
-            // Oh boy a recursive instantiation! The recursive arguments here
-            // are pretty simple, and the only slightly-meaty one is how
-            // arguments are pulled from `args` and pushed directly into the
-            // builder specified, which should be an easy enough
-            // copy-the-pointer operation in all cases.
+            // Oh boy a recursive instantiation!
             //
-            // Note that this recursive call shouldn't result in an infinite
-            // loop because of wasm module validation which requires everything
-            // to be a DAG. Additionally the recursion should also be bounded
-            // due to validation. We may one day need to make this an iterative
-            // loop, however.
+            // We use our local index space of modules to find the module to
+            // instantiate and argument lookup is defined as looking inside of
+            // `args`. Like above with aliases all type checking is eschewed
+            // because only valid modules should reach this point.
             //
-            // Also note that there's some unsafety here around cloning
-            // `InstanceHandle` because the handle may not live long enough, but
-            // we're doing all of this in the context of our `Store` argument
-            // above so we should be safe here.
+            // Note that it's thought that due to the acyclic nature of
+            // instantiation this can't loop to blow the native stack, and
+            // validation should in theory ensure this has a bounded depth.
+            // Despite this we may need to change this to a loop instead of
+            // recursion one day.
             Initializer::Instantiate { module, args } => {
                 let handle = instantiate(
                     store,
@@ -134,6 +115,33 @@ fn instantiate(
                 )?;
                 imports.instances.push(handle);
             }
+
+            // A new module is being defined, and the source of this module is
+            // our module's list of closed-over-modules.
+            //
+            // This is used for outer aliases.
+            Initializer::DefineModule(upvar_index) => {
+                imports
+                    .modules
+                    .push(module.module_upvar(*upvar_index).clone());
+            }
+
+            // A new module is defined, created from a set of compiled
+            // artifacts. The new module value will be created with the
+            // specified artifacts being closed over as well as the specified
+            // set of module values in our index/upvar index spaces being closed
+            // over.
+            //
+            // This is used for defining submodules.
+            Initializer::CreateModule {
+                artifact_index,
+                artifacts,
+                modules,
+            } => {
+                let submodule =
+                    module.create_submodule(*artifact_index, artifacts, modules, &imports.modules);
+                imports.modules.push(submodule);
+            }
         }
     }