module-linking: Implement outer module aliases

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);
+            }
         }
     }