From 76e9fd12dc664e5a61c6b368961d66bc86c49837 Mon Sep 17 00:00:00 2001 From: Sy Brand Date: Tue, 13 Jan 2026 16:05:37 +0000 Subject: [PATCH 1/4] Add task init wrapper --- crates/wit-component/src/encoding.rs | 259 ++++++++++++++++++++++--- crates/wit-component/src/validation.rs | 38 ++++ 2 files changed, 268 insertions(+), 29 deletions(-) diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs index 15d82b23f8..c427463283 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,13 @@ pub struct EncodingState<'a> { /// Metadata about the world inferred from the input to `ComponentEncoder`. info: &'a ComponentWorld<'a>, + + /// Maps from original export name to wrapper function index. + /// Used to wrap exports with __wasilibc_init_task calls. + export_wrappers: HashMap, + + /// Index of wasilibc init task wrapper module instance, if any. + wrapper_instance_index: Option, } impl<'a> EncodingState<'a> { @@ -600,6 +633,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 +725,9 @@ impl<'a> EncodingState<'a> { | Export::GeneralPurposeImportRealloc | Export::Initialize | Export::ReallocForAdapter - | Export::IndirectFunctionTable => continue, + | Export::IndirectFunctionTable + | Export::WasiLibCInitTask + | Export::WasiLibCInitTaskAsync => continue, } } @@ -1014,32 +1053,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 +1066,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_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 +2207,187 @@ 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 mut funcs_to_wrap = Vec::new(); + for (core_name, export) in exports.iter() { + match export { + Export::WorldFunc(key, _, abi) => { + let func = match &world.exports[key] { + WorldItem::Function(f) => f, + _ => continue, + }; + funcs_to_wrap.push((core_name, func, abi)); + } + Export::InterfaceFunc(_, id, func_name, abi) => { + let func = &resolve.interfaces[*id].functions[func_name.as_str()]; + funcs_to_wrap.push((core_name, func, abi)); + } + _ => continue, + } + } + + 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; + + 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); + + 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 functions + 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))], + ); + + self.wrapper_instance_index = Some(wrapper_instance); + + // 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_wrappers.insert(name.into(), wrapper_idx); + } + + Ok(()) + } } /// A list of "shims" which start out during the component instantiation process @@ -3123,6 +3322,8 @@ impl ComponentEncoder { exported_instances: Default::default(), aliased_core_items: Default::default(), info: &world, + export_wrappers: HashMap::new(), + wrapper_instance_index: None, }; 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]") } From 2c50cf3570be745b7a8680566a3fcbce6b3252ec Mon Sep 17 00:00:00 2001 From: Sy Brand Date: Tue, 13 Jan 2026 16:08:40 +0000 Subject: [PATCH 2/4] Cleanup --- crates/wit-component/src/encoding.rs | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs index c427463283..389403386e 100644 --- a/crates/wit-component/src/encoding.rs +++ b/crates/wit-component/src/encoding.rs @@ -421,12 +421,9 @@ pub struct EncodingState<'a> { /// Metadata about the world inferred from the input to `ComponentEncoder`. info: &'a ComponentWorld<'a>, - /// Maps from original export name to wrapper function index. + /// Maps from original export name to task initialization wrapper function index. /// Used to wrap exports with __wasilibc_init_task calls. - export_wrappers: HashMap, - - /// Index of wasilibc init task wrapper module instance, if any. - wrapper_instance_index: Option, + export_task_initialization_wrappers: HashMap, } impl<'a> EncodingState<'a> { @@ -1068,7 +1065,7 @@ impl<'a> EncodingState<'a> { let instance_index = self.instance_for(module); // 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_wrappers.get(core_name) { + 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) @@ -2377,13 +2374,11 @@ impl<'a> EncodingState<'a> { [("", ModuleArg::Instance(wrapper_args_idx))], ); - self.wrapper_instance_index = Some(wrapper_instance); - // 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_wrappers.insert(name.into(), wrapper_idx); + self.export_task_initialization_wrappers.insert(name.into(), wrapper_idx); } Ok(()) @@ -3322,8 +3317,7 @@ impl ComponentEncoder { exported_instances: Default::default(), aliased_core_items: Default::default(), info: &world, - export_wrappers: HashMap::new(), - wrapper_instance_index: None, + export_task_initialization_wrappers: HashMap::new(), }; state.encode_imports(&self.import_name_map)?; state.encode_core_modules(); From 9e2b0bfc929c7247c58bba865ecc2c9ea043c7ad Mon Sep 17 00:00:00 2001 From: Sy Brand Date: Tue, 13 Jan 2026 16:12:59 +0000 Subject: [PATCH 3/4] Cleanup --- crates/wit-component/src/encoding.rs | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs index 389403386e..687867cab0 100644 --- a/crates/wit-component/src/encoding.rs +++ b/crates/wit-component/src/encoding.rs @@ -2232,23 +2232,20 @@ impl<'a> EncodingState<'a> { // Collect the exports that we will need to wrap, alongside information // that we'll need to build the wrappers. - let mut funcs_to_wrap = Vec::new(); - for (core_name, export) in exports.iter() { - match export { - Export::WorldFunc(key, _, abi) => { - let func = match &world.exports[key] { - WorldItem::Function(f) => f, - _ => continue, - }; - funcs_to_wrap.push((core_name, func, abi)); - } + 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()]; - funcs_to_wrap.push((core_name, func, abi)); + Some((core_name, func, abi)) } - _ => continue, - } - } + _ => None, + }) + .collect(); if funcs_to_wrap.is_empty() { // No exports, so no wrappers are needed. From 8fb288728e00152932035e70a6c065f638f5f304 Mon Sep 17 00:00:00 2001 From: Sy Brand Date: Tue, 13 Jan 2026 16:15:26 +0000 Subject: [PATCH 4/4] Cleanup --- crates/wit-component/src/encoding.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/wit-component/src/encoding.rs b/crates/wit-component/src/encoding.rs index 687867cab0..97187098a1 100644 --- a/crates/wit-component/src/encoding.rs +++ b/crates/wit-component/src/encoding.rs @@ -2263,6 +2263,7 @@ impl<'a> EncodingState<'a> { 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, @@ -2332,8 +2333,8 @@ impl<'a> EncodingState<'a> { .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, @@ -2349,7 +2350,7 @@ impl<'a> EncodingState<'a> { 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 functions + // 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);