diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs index 15d82b23f8..97187098a1 100644 --- a/crates/wit-component/src/encoding.rs +++ b/crates/wit-component/src/encoding.rs @@ -116,6 +116,32 @@ fn to_val_type(ty: &WasmType) -> ValType { } } + fn import_func_name(f: &Function) -> String { + match f.kind { + FunctionKind::Freestanding | FunctionKind::AsyncFreestanding => { + format!("import-func-{}", f.item_name()) + } + + // transform `[method]foo.bar` into `import-method-foo-bar` to + // have it be a valid kebab-name which can't conflict with + // anything else. + // + // There's probably a better and more "formal" way to do this + // but quick-and-dirty string manipulation should work well + // enough for now hopefully. + FunctionKind::Method(_) + | FunctionKind::AsyncMethod(_) + | FunctionKind::Static(_) + | FunctionKind::AsyncStatic(_) + | FunctionKind::Constructor(_) => { + format!( + "import-{}", + f.name.replace('[', "").replace([']', '.', ' '], "-") + ) + } + } + } + bitflags::bitflags! { /// Options in the `canon lower` or `canon lift` required for a particular /// function. @@ -394,6 +420,10 @@ pub struct EncodingState<'a> { /// Metadata about the world inferred from the input to `ComponentEncoder`. info: &'a ComponentWorld<'a>, + + /// Maps from original export name to task initialization wrapper function index. + /// Used to wrap exports with __wasilibc_init_task calls. + export_task_initialization_wrappers: HashMap, } impl<'a> EncodingState<'a> { @@ -600,6 +630,10 @@ impl<'a> EncodingState<'a> { // at the end. self.instantiate_main_module(&shims)?; + // Create any wrappers needed for initializing tasks if wasilibc is + // being used. + self.create_export_task_initialization_wrappers()?; + // Separate the adapters according which should be instantiated before // and after indirect lowerings are encoded. let (before, after) = self @@ -688,7 +722,9 @@ impl<'a> EncodingState<'a> { | Export::GeneralPurposeImportRealloc | Export::Initialize | Export::ReallocForAdapter - | Export::IndirectFunctionTable => continue, + | Export::IndirectFunctionTable + | Export::WasiLibCInitTask + | Export::WasiLibCInitTaskAsync => continue, } } @@ -1014,32 +1050,6 @@ impl<'a> EncodingState<'a> { name } } - - fn import_func_name(f: &Function) -> String { - match f.kind { - FunctionKind::Freestanding | FunctionKind::AsyncFreestanding => { - format!("import-func-{}", f.item_name()) - } - - // transform `[method]foo.bar` into `import-method-foo-bar` to - // have it be a valid kebab-name which can't conflict with - // anything else. - // - // There's probably a better and more "formal" way to do this - // but quick-and-dirty string manipulation should work well - // enough for now hopefully. - FunctionKind::Method(_) - | FunctionKind::AsyncMethod(_) - | FunctionKind::Static(_) - | FunctionKind::AsyncStatic(_) - | FunctionKind::Constructor(_) => { - format!( - "import-{}", - f.name.replace('[', "").replace([']', '.', ' '], "-") - ) - } - } - } } fn encode_lift( @@ -1053,8 +1063,13 @@ impl<'a> EncodingState<'a> { let resolve = &self.info.encoder.metadata.resolve; let metadata = self.info.module_metadata_for(module); let instance_index = self.instance_for(module); - let core_func_index = - self.core_alias_export(Some(core_name), instance_index, core_name, ExportKind::Func); + // If we generated a wasilibc init task wrapper for this export, use that, + // otherwise alias the original export. + let core_func_index = if let Some(&wrapper_idx) = self.export_task_initialization_wrappers.get(core_name) { + wrapper_idx + } else { + self.core_alias_export(Some(core_name), instance_index, core_name, ExportKind::Func) + }; let exports = self.info.exports_for(module); let options = RequiredOptions::for_export( @@ -2189,6 +2204,183 @@ impl<'a> EncodingState<'a> { .core_alias_export(debug_name, instance, name, kind) }) } + + /// wasi-libc defines `__wasilibc_init_task(_async)` functions that must be called + /// at the start of every exported function to set up the stack pointer and + /// thread-local storage. To achieve this, we create a wrapper module called + /// `wasilibc-init-wrappers` that imports the original exports and the + /// task initialization functions, and defines wrapper functions that call + /// the relevant task initialization function before delegating to the original export. + /// We then instantiate this wrapper module and use its exports as the final + /// exports of the component. If we don't find a `__wasilibc_init_task` export, + /// we elide the wrapper module entirely. + fn create_export_task_initialization_wrappers(&mut self) -> Result<()> { + let instance_index = self.instance_index.unwrap(); + let resolve = &self.info.encoder.metadata.resolve; + let world = &resolve.worlds[self.info.encoder.metadata.world]; + let exports = self.info.exports_for(CustomModule::Main); + let wasilibc_init_task = "__wasilibc_init_task"; + let wasilibc_init_task_async = "__wasilibc_init_task_async"; + + let wasilibc_init_task_export = exports.wasilibc_init_task(); + let wasilibc_init_task_async_export = exports.wasilibc_init_task_async(); + if wasilibc_init_task_export.is_none() || wasilibc_init_task_async_export.is_none() { + // __wasilibc_init_task(_async) was not exported by the main module, + // so no wrappers are needed. + return Ok(()); + } + + // Collect the exports that we will need to wrap, alongside information + // that we'll need to build the wrappers. + let funcs_to_wrap: Vec<_> = exports + .iter() + .flat_map(|(core_name, export)| match export { + Export::WorldFunc(key, _, abi) => match &world.exports[key] { + WorldItem::Function(f) => Some((core_name, f, abi)), + _ => None, + }, + Export::InterfaceFunc(_, id, func_name, abi) => { + let func = &resolve.interfaces[*id].functions[func_name.as_str()]; + Some((core_name, func, abi)) + } + _ => None, + }) + .collect(); + + if funcs_to_wrap.is_empty() { + // No exports, so no wrappers are needed. + return Ok(()); + } + + // Now we build the wrapper module + let mut types = TypeSection::new(); + let mut imports = ImportSection::new(); + let mut functions = FunctionSection::new(); + let mut exports_section = ExportSection::new(); + let mut code = CodeSection::new(); + + // Type for __wasilibc_init_task(_async): () -> () + types.ty().function([], []); + let wasilibc_init_task_type_idx = 0; + + // Import __wasilibc_init_task and __wasilibc_init_task_async into the wrapper module + imports.import( + "", + wasilibc_init_task, + EntityType::Function(wasilibc_init_task_type_idx), + ); + imports.import( + "", + wasilibc_init_task_async, + EntityType::Function(wasilibc_init_task_type_idx), + ); + let wasilibc_init_task_func_idx = 0u32; + let wasilibc_init_task_async_func_idx = 1u32; + + let mut type_indices = HashMap::new(); + let mut next_type_idx = 1u32; + let mut next_func_idx = 2u32; + + // Create wrapper functions for each export + for &(name, func, abi) in funcs_to_wrap.iter() { + let sig = resolve.wasm_signature(*abi, func); + let type_idx = *type_indices.entry(sig.clone()).or_insert_with(|| { + let idx = next_type_idx; + types.ty().function( + sig.params.iter().map(to_val_type), + sig.results.iter().map(to_val_type), + ); + next_type_idx += 1; + idx + }); + + imports.import( + "", + &import_func_name(func), + EntityType::Function(type_idx), + ); + let orig_func_idx = next_func_idx; + next_func_idx += 1; + + let wrapper_func_idx = next_func_idx; + functions.function(type_idx); + + let mut func = wasm_encoder::Function::new([]); + if abi.is_async() { + func.instruction(&Instruction::Call(wasilibc_init_task_async_func_idx)); + } else { + func.instruction(&Instruction::Call(wasilibc_init_task_func_idx)); + } + for i in 0..sig.params.len() as u32 { + func.instruction(&Instruction::LocalGet(i)); + } + func.instruction(&Instruction::Call(orig_func_idx)); + func.instruction(&Instruction::End); + code.function(&func); + + exports_section.export(name, ExportKind::Func, wrapper_func_idx); + next_func_idx += 1; + } + + let mut wrapper_module = Module::new(); + wrapper_module.section(&types); + wrapper_module.section(&imports); + wrapper_module.section(&functions); + wrapper_module.section(&exports_section); + wrapper_module.section(&code); + + let wrapper_module_idx = self + .component + .core_module(Some("wasilibc-init-wrappers"), &wrapper_module); + + // Prepare imports for instantiating the wrapper module + let mut wrapper_imports = Vec::new(); + let init_idx = self.core_alias_export( + Some(wasilibc_init_task), + instance_index, + wasilibc_init_task, + ExportKind::Func, + ); + let init_async_idx = self.core_alias_export( + Some(wasilibc_init_task_async), + instance_index, + wasilibc_init_task_async, + ExportKind::Func, + ); + wrapper_imports.push((wasilibc_init_task.into(), ExportKind::Func, init_idx)); + wrapper_imports.push((wasilibc_init_task_async.into(), ExportKind::Func, init_async_idx)); + + // Import all original exports to be wrapped + for (name, func, _) in &funcs_to_wrap { + let orig_idx = + self.core_alias_export(Some(name), instance_index, name, ExportKind::Func); + wrapper_imports.push(( + import_func_name(func), + ExportKind::Func, + orig_idx, + )); + } + + let wrapper_args_idx = self.component.core_instantiate_exports( + Some("wasilibc-init-wrappers-args"), + wrapper_imports.iter().map(|(n, k, i)| (n.as_str(), *k, *i)), + ); + + let wrapper_instance = self.component.core_instantiate( + Some("wasilibc-init-wrappers-instance"), + wrapper_module_idx, + [("", ModuleArg::Instance(wrapper_args_idx))], + ); + + // Map original names to wrapper indices + for (name, _, _) in funcs_to_wrap { + let wrapper_idx = + self.core_alias_export(Some(&name), wrapper_instance, &name, ExportKind::Func); + self.export_task_initialization_wrappers.insert(name.into(), wrapper_idx); + } + + Ok(()) + } } /// A list of "shims" which start out during the component instantiation process @@ -3123,6 +3315,7 @@ impl ComponentEncoder { exported_instances: Default::default(), aliased_core_items: Default::default(), info: &world, + export_task_initialization_wrappers: HashMap::new(), }; state.encode_imports(&self.import_name_map)?; state.encode_core_modules(); diff --git a/crates/wit-component/src/validation.rs b/crates/wit-component/src/validation.rs index 88d7f80833..f7cb3a6d98 100644 --- a/crates/wit-component/src/validation.rs +++ b/crates/wit-component/src/validation.rs @@ -1111,6 +1111,12 @@ pub enum Export { /// __indirect_function_table, used for `thread.new-indirect` IndirectFunctionTable, + + /// __wasilibc_init_task, used for wasi-libc initialization + WasiLibCInitTask, + + /// __wasilibc_init_task_async, used for wasi-libc initialization for async-lifted exports + WasiLibCInitTaskAsync, } impl ExportMap { @@ -1209,6 +1215,14 @@ impl ExportMap { let expected = FuncType::new([], []); validate_func_sig(name, &expected, ty)?; return Ok(Some(Export::Initialize)); + } else if Some(name) == names.export_wasilibc_init_task() { + let expected = FuncType::new([], []); + validate_func_sig(name, &expected, ty)?; + return Ok(Some(Export::WasiLibCInitTask)); + } else if Some(name) == names.export_wasilibc_init_task_async() { + let expected = FuncType::new([], []); + validate_func_sig(name, &expected, ty)?; + return Ok(Some(Export::WasiLibCInitTaskAsync)); } let full_name = name; @@ -1372,6 +1386,16 @@ impl ExportMap { self.find(|t| matches!(t, Export::IndirectFunctionTable)) } + /// Returns the `__wasilibc_init_task` function, if exported, for this module. + pub fn wasilibc_init_task(&self) -> Option<&str> { + self.find(|t| matches!(t, Export::WasiLibCInitTask)) + } + + /// Returns the `__wasilibc_init_task_async` function, if exported, for this module. + pub fn wasilibc_init_task_async(&self) -> Option<&str> { + self.find(|t| matches!(t, Export::WasiLibCInitTaskAsync)) + } + /// Returns the `_initialize` intrinsic, if exported, for this module. pub fn initialize(&self) -> Option<&str> { self.find(|m| matches!(m, Export::Initialize)) @@ -1529,6 +1553,8 @@ trait NameMangling { fn export_initialize(&self) -> &str; fn export_realloc(&self) -> &str; fn export_indirect_function_table(&self) -> Option<&str>; + fn export_wasilibc_init_task(&self) -> Option<&str>; + fn export_wasilibc_init_task_async(&self) -> Option<&str>; fn resource_drop_name<'a>(&self, name: &'a str) -> Option<&'a str>; fn resource_new_name<'a>(&self, name: &'a str) -> Option<&'a str>; fn resource_rep_name<'a>(&self, name: &'a str) -> Option<&'a str>; @@ -1673,6 +1699,12 @@ impl NameMangling for Standard { fn export_indirect_function_table(&self) -> Option<&str> { None } + fn export_wasilibc_init_task(&self) -> Option<&str> { + None + } + fn export_wasilibc_init_task_async(&self) -> Option<&str> { + None + } fn resource_drop_name<'a>(&self, name: &'a str) -> Option<&'a str> { name.strip_suffix("_drop") } @@ -2113,6 +2145,12 @@ impl NameMangling for Legacy { fn export_indirect_function_table(&self) -> Option<&str> { Some("__indirect_function_table") } + fn export_wasilibc_init_task(&self) -> Option<&str> { + Some("__wasilibc_init_task") + } + fn export_wasilibc_init_task_async(&self) -> Option<&str> { + Some("__wasilibc_init_task_async") + } fn resource_drop_name<'a>(&self, name: &'a str) -> Option<&'a str> { name.strip_prefix("[resource-drop]") }