diff --git a/crates/wasmlink-cli/src/lib.rs b/crates/wasmlink-cli/src/lib.rs index e89dc3c47..a1d8d6c89 100644 --- a/crates/wasmlink-cli/src/lib.rs +++ b/crates/wasmlink-cli/src/lib.rs @@ -41,13 +41,13 @@ pub struct App { #[structopt(long, short = "p", value_name = "PROFILE")] pub profile: String, - /// The path of the output linked module; defaults to replacing the main module. + /// The path of the output linked module; defaults to replacing the given module. #[structopt(long, short = "o", value_name = "OUTPUT", parse(from_os_str))] pub output: Option, - /// The main module to link. + /// The module to link. #[structopt(index = 1, value_name = "MODULE", parse(from_os_str))] - pub main: PathBuf, + pub module: PathBuf, } impl App { @@ -57,15 +57,15 @@ impl App { bail!("at least one import module must be specified"); } - let main_bytes = wat::parse_file(&self.main) - .with_context(|| format!("failed to parse main module `{}`", self.main.display()))?; + let module_bytes = wat::parse_file(&self.module) + .with_context(|| format!("failed to parse module `{}`", self.module.display()))?; - let main_module = Module::new( - self.main.file_name().unwrap().to_str().unwrap(), - &main_bytes, + let module = Module::new( + self.module.file_name().unwrap().to_str().unwrap(), + &module_bytes, [], ) - .with_context(|| format!("failed to parse main module `{}`", self.main.display()))?; + .with_context(|| format!("failed to parse module `{}`", self.module.display()))?; let import_bytes = self .modules @@ -118,8 +118,8 @@ impl App { let linker = Linker::new(Profile::new()); - let output = self.output.as_ref().unwrap_or(&self.main); - std::fs::write(output, linker.link(&main_module, &import_modules)?) + let output = self.output.as_ref().unwrap_or(&self.module); + std::fs::write(output, linker.link(&module, &import_modules)?) .with_context(|| format!("failed to write to output module `{}`", output.display()))?; Ok(()) diff --git a/crates/wasmlink/src/linker.rs b/crates/wasmlink/src/linker.rs index d70b2e80e..1263c4faa 100644 --- a/crates/wasmlink/src/linker.rs +++ b/crates/wasmlink/src/linker.rs @@ -37,6 +37,7 @@ struct LinkedModule<'a> { module_map: HashMap<&'a ModuleAdapter<'a>, (u32, Option)>, instances: Vec<(u32, Vec<(&'a str, u32)>)>, func_aliases: Vec<(u32, &'a str)>, + memory_aliases: Vec<(u32, &'a str)>, table_aliases: Vec<(u32, &'a str)>, segments: Vec<(u32, Vec)>, exports: Vec<(&'a str, wasm_encoder::Export)>, @@ -126,16 +127,32 @@ impl<'a> LinkedModule<'a> { linked.instances.push((index, args)); } - // Instantiate the main module - let (main_index, _) = linked.instantiate(graph, NodeIndex::new(0), None)?; - - // Re-export the start function - let start_index = linked.func_aliases.len() as u32; - linked.func_aliases.push((main_index, "_start")); - linked.exports.push(( - "_start", - wasm_encoder::Export::Function(linked.imports.len() as u32 + start_index), - )); + // Instantiate the root module + let (root_index, _) = linked.instantiate(graph, NodeIndex::new(0), None)?; + + let root = &graph[NodeIndex::new(0)]; + + // Re-export all supported exports from the root module + for export in &root.module.exports { + match export.kind { + ExternalKind::Function => { + let func_index = linked.func_aliases.len() as u32; + linked.func_aliases.push((root_index, export.field)); + linked.exports.push(( + export.field, + wasm_encoder::Export::Function(linked.imports.len() as u32 + func_index), + )); + } + ExternalKind::Memory => { + let memory_index = linked.memory_aliases.len() as u32; + linked.memory_aliases.push((root_index, export.field)); + linked + .exports + .push((export.field, wasm_encoder::Export::Memory(memory_index))); + } + _ => {} + } + } Ok(linked) } @@ -326,6 +343,10 @@ impl<'a> LinkedModule<'a> { section.instance_export(*index, wasm_encoder::ItemKind::Table, name); } + for (index, name) in &self.memory_aliases { + section.instance_export(*index, wasm_encoder::ItemKind::Memory, name); + } + module.section(§ion); } @@ -370,21 +391,8 @@ impl Linker { /// Links the given module with the given set of imported modules. /// /// On success, returns a vector of bytes representing the linked module. - pub fn link(&self, main: &Module, imports: &HashMap<&str, Module>) -> Result> { - if !main.exports.iter().any(|e| match (e.field, e.kind) { - ("_start", ExternalKind::Function) => { - let ty = main.func_type(e.index).unwrap(); - ty.params.is_empty() && ty.returns.is_empty() - } - _ => false, - }) { - bail!( - "main module `{}` must export a start function that has no parameters or results", - main.name - ); - } - - let (graph, needs_runtime) = self.build_graph(main, imports)?; + pub fn link(&self, module: &Module, imports: &HashMap<&str, Module>) -> Result> { + let (graph, needs_runtime) = self.build_graph(module, imports)?; let module = LinkedModule::new(&graph, needs_runtime, &self.profile)?; @@ -393,16 +401,16 @@ impl Linker { fn build_graph<'a>( &self, - main: &'a Module, + module: &'a Module, imports: &'a HashMap<&str, Module>, ) -> Result<(Graph, ()>, bool)> { let mut queue: Vec<(Option, &Module)> = Vec::new(); let mut seen = HashMap::new(); let mut graph: Graph = Graph::new(); - let mut needs_runtime = main.has_resources; + let mut needs_runtime = module.has_resources; - queue.push((None, main)); + queue.push((None, module)); let mut next_resource_id = 0; @@ -476,36 +484,6 @@ impl Linker { mod test { use super::*; - #[test] - fn it_errors_on_missing_main_function() -> Result<()> { - let bytes = wat::parse_str(r#"(module)"#)?; - let main = Module::new("main", &bytes, [])?; - - let linker = Linker::new(Profile::new()); - - assert_eq!( - linker.link(&main, &HashMap::new()).unwrap_err().to_string(), - "main module `main` must export a start function that has no parameters or results" - ); - - Ok(()) - } - - #[test] - fn it_errors_on_incorrect_main_function() -> Result<()> { - let bytes = wat::parse_str(r#"(module (func (export "_start") (param i32)))"#)?; - let main = Module::new("main", &bytes, [])?; - - let linker = Linker::new(Profile::new()); - - assert_eq!( - linker.link(&main, &HashMap::new()).unwrap_err().to_string(), - "main module `main` must export a start function that has no parameters or results" - ); - - Ok(()) - } - #[test] fn it_errors_on_missing_import() -> Result<()> { let bytes = wat::parse_str( @@ -743,7 +721,7 @@ mod test { #[test] fn it_links_with_interface() -> Result<()> { let bytes = wat::parse_str( - r#"(module (import "a" "a" (func (param i32 i32))) (func (export "_start")) (func (export "canonical_abi_realloc") (param i32 i32 i32 i32) (result i32) unreachable))"#, + r#"(module (import "a" "a" (func (param i32 i32))) (func (export "_start")) (func (export "canonical_abi_realloc") (param i32 i32 i32 i32) (result i32) unreachable) (memory (export "memory") 0))"#, )?; let a = wat::parse_str( r#"(module (import "wasi_snapshot_preview1" "a" (func)) (func (export "a") (param i32 i32)) (memory (export "memory") 0) (func (export "canonical_abi_realloc") (param i32 i32 i32 i32) (result i32) unreachable) (func (export "canonical_abi_free") (param i32 i32 i32)))"#, @@ -777,8 +755,10 @@ mod test { (func (;1;) (type 1)) (func (;2;) (type 2) (param i32 i32 i32 i32) (result i32) unreachable) + (memory (;0;) 0) (export \"_start\" (func 1)) - (export \"canonical_abi_realloc\" (func 2))) + (export \"canonical_abi_realloc\" (func 2)) + (export \"memory\" (memory 0))) (module (;1;) (type (;0;) (func)) (type (;1;) (func (param i32 i32))) @@ -853,8 +833,12 @@ mod test { (import \"$parent\" (instance 2)))) (alias 3 \"a\" (func (;1;))) (alias 2 \"_start\" (func (;2;))) + (alias 2 \"canonical_abi_realloc\" (func (;3;))) (alias 1 \"$funcs\" (table (;0;))) + (alias 2 \"memory\" (memory (;0;))) (export \"_start\" (func 2)) + (export \"canonical_abi_realloc\" (func 3)) + (export \"memory\" (memory 0)) (elem (;0;) (i32.const 0) funcref (ref.func 1)))" );