From 07820828b5a0d066e395576321f6658f6a5f0278 Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 12:46:48 -0700 Subject: [PATCH 1/2] refactor wrapped --- src/ast.rs | 441 +++++++++++++++++++----------------------- src/codegen/mod.rs | 13 +- src/codegen/record.rs | 6 +- src/instrument.rs | 3 +- src/traits/proxy.rs | 5 +- src/util.rs | 32 ++- 6 files changed, 241 insertions(+), 259 deletions(-) diff --git a/src/ast.rs b/src/ast.rs index cbdcf48..d7bba58 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -1,7 +1,7 @@ use crate::Mode; use crate::instrument::InstrumentArgs; use crate::util::*; -use anyhow::Result; +use anyhow::{Result, bail}; use std::collections::{BTreeMap, BTreeSet}; use std::path::Path; use wit_bindgen_core::{Files, Source}; @@ -11,16 +11,16 @@ use wit_parser::*; pub struct Opt<'a> { args: &'a InstrumentArgs, mode: Mode, - imports: LinkInfo, - exports: LinkInfo, - main: LinkInfo, + /// How each import of the main component is linked. + main_imports: BTreeMap, + main_exports: BTreeSet, + /// How each import of the generated `imports` component is linked. + imports_links: BTreeMap, + /// How each import of the generated `exports` component is linked. + exports_links: BTreeMap, need_debug: bool, } -#[derive(Default)] -struct LinkInfo { - imports: BTreeMap, - exports: BTreeSet, -} +/// The instance in the composed component that provides an import. enum LinkType { Debug, Recorder, @@ -35,124 +35,97 @@ impl<'a> Opt<'a> { Self { args, mode, - imports: LinkInfo::default(), - exports: LinkInfo::default(), - main: LinkInfo::default(), + main_imports: BTreeMap::new(), + main_exports: BTreeSet::new(), + imports_links: BTreeMap::new(), + exports_links: BTreeMap::new(), need_debug: false, } } - fn generate_main_wit(&mut self, resolve: &Resolve, id: WorldId, files: &mut Files) { - let mut out = Source::default(); - let world = &resolve.worlds[id]; - let recorder = "proxy:recorder/"; - out.push_str("package component:proxy;\n"); - out.push_str("world imports {\n"); + fn push_mode_imports(&self, out: &mut Source) { match self.mode { Mode::Record | Mode::Replay => { out.push_str(&format!( - "import {recorder}{}@0.1.0;\n", + "import proxy:recorder/{}@0.1.0;\n", ident(self.mode.to_str()) )); } - Mode::Fuzz => { - out.push_str(&format!("import proxy:util/debug;\n")); - } - Mode::Dialog => { - out.push_str(&format!("import proxy:util/dialog;\n")); - } - }; + Mode::Fuzz => out.push_str("import proxy:util/debug;\n"), + Mode::Dialog => out.push_str("import proxy:util/dialog;\n"), + } + } + fn generate_main_wit( + &mut self, + resolve: &Resolve, + id: WorldId, + files: &mut Files, + ) -> Result<()> { + let mut out = Source::default(); + let world = &resolve.worlds[id]; + out.push_str("package component:proxy;\n"); + out.push_str("world imports {\n"); + self.push_mode_imports(&mut out); out.push_str("export proxy:conversion/conversion;\n"); - for (name, import) in &world.imports { - match import { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - // Don't virtualize util imports - if name.starts_with("proxy:util/") { - self.need_debug = true; - self.main.imports.insert(name.clone(), LinkType::Debug); - out.push_str(&format!("import {name};\n")); - continue; - } - self.main.imports.insert(name.clone(), LinkType::Imports); - match self.mode { - Mode::Record => { - out.push_str(&format!("import {name};\n")); - out.push_str(&format!("export wrapped-{name};\n")); - } - Mode::Replay | Mode::Fuzz | Mode::Dialog => { - out.push_str(&format!("export {name};\n")) - } - } + for name in interface_names(resolve, &world.imports)? { + // Don't virtualize util imports + if name.starts_with("proxy:util/") { + self.need_debug = true; + out.push_str(&format!("import {name};\n")); + self.main_imports.insert(name, LinkType::Debug); + continue; + } + match self.mode { + Mode::Record => { + out.push_str(&format!("import {name};\n")); + out.push_str(&format!("export {};\n", wrapped_name(&name))); + } + Mode::Replay | Mode::Fuzz | Mode::Dialog => { + out.push_str(&format!("export {name};\n")) } - _ => todo!(), } + self.main_imports.insert(name, LinkType::Imports); } out.push_str("}\n"); out.push_str("world tmp-exports {\n"); - for (name, export) in &world.exports { - match export { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - self.main.exports.insert(name.clone()); - match self.mode { - Mode::Record => { - out.push_str(&format!("import wrapped-{name};\n")); - out.push_str(&format!("export {name};\n")); - } - Mode::Replay | Mode::Fuzz | Mode::Dialog => { - out.push_str(&format!("import {name};\n")); - } - } + for name in interface_names(resolve, &world.exports)? { + match self.mode { + Mode::Record => { + out.push_str(&format!("import {};\n", wrapped_name(&name))); + out.push_str(&format!("export {name};\n")); + } + Mode::Replay | Mode::Fuzz | Mode::Dialog => { + out.push_str(&format!("import {name};\n")); } - _ => todo!(), } + self.main_exports.insert(name); } if matches!(self.mode, Mode::Replay | Mode::Fuzz | Mode::Dialog) { out.push_str("export proxy:recorder/start-replay@0.1.0;\n") } out.push_str("}\n"); files.push("component.wit", out.as_bytes()); + Ok(()) } - pub fn generate_exports_world(&self, resolve: &Resolve, id: WorldId, files: &mut Files) { + pub fn generate_exports_world( + &self, + resolve: &Resolve, + id: WorldId, + files: &mut Files, + ) -> Result<()> { let mut out = Source::default(); let world = &resolve.worlds[id]; - let recorder = "proxy:recorder/"; out.push_str("world exports {\n"); - match self.mode { - Mode::Record | Mode::Replay => { - out.push_str(&format!( - "import {recorder}{}@0.1.0;\n", - ident(self.mode.to_str()) - )); - } - Mode::Fuzz => { - out.push_str(&format!("import proxy:util/debug;\n")); - } - Mode::Dialog => { - out.push_str(&format!("import proxy:util/dialog;\n")); - } - }; + self.push_mode_imports(&mut out); out.push_str("import proxy:conversion/conversion;\n"); - for (name, import) in &world.imports { - match import { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - out.push_str(&format!("import {name};\n")); - } - _ => todo!(), - } + for name in interface_names(resolve, &world.imports)? { + out.push_str(&format!("import {name};\n")); } - for (name, export) in &world.exports { - match export { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - out.push_str(&format!("export {name};\n")); - } - _ => todo!(), - } + for name in interface_names(resolve, &world.exports)? { + out.push_str(&format!("export {name};\n")); } out.push_str("}\n"); files.push("component.wit", out.as_bytes()); + Ok(()) } pub fn generate_wac( &mut self, @@ -160,7 +133,6 @@ impl<'a> Opt<'a> { exports_wasm: &Path, out_dir: &Path, ) -> Result<()> { - // Get WIT from wasm to account for unused imports being optimized away self.load_imports(imports_wasm)?; self.load_exports(exports_wasm)?; let mut out = Source::default(); @@ -173,41 +145,20 @@ impl<'a> Opt<'a> { } out.push_str("let imports = new import:proxy {\n"); let mut has_host = false; - for (name, link_type) in &self.imports.imports { - match link_type { - LinkType::Debug => { - out.push_str(&format!("\"{name}\": debug[\"{name}\"] ,\n")); - } - LinkType::Recorder => { - if self.args.use_host_recorder { - has_host = true; - } else { - out.push_str(&format!("\"{name}\": recorder[\"{name}\"] ,\n")); - } - } - LinkType::Host => { - has_host = true; - } - LinkType::Imports | LinkType::Main => unreachable!(), - } - } - if has_host { - out.push_str("...,\n"); + for (name, link_type) in &self.imports_links { + has_host |= self.push_external_link(&mut out, name, link_type); } - out.push_str("};\n"); + close_instance(&mut out, has_host); out.push_str("let main = new root:component {\n"); - for (name, link_type) in &self.main.imports { + for (name, link_type) in &self.main_imports { match link_type { - LinkType::Debug => { - out.push_str(&format!("\"{name}\": debug[\"{name}\"] ,\n")); - } + LinkType::Debug => push_link(&mut out, name, "debug", name), LinkType::Imports => { - let prefix = if matches!(self.mode, Mode::Record) { - "wrapped-" - } else { - "" + let export = match self.mode { + Mode::Record => wrapped_name(name), + _ => name.clone(), }; - out.push_str(&format!("\"{name}\": imports[\"{prefix}{name}\"] ,\n")); + push_link(&mut out, name, "imports", &export); } LinkType::Host | LinkType::Main | LinkType::Recorder => unreachable!(), } @@ -215,48 +166,45 @@ impl<'a> Opt<'a> { out.push_str("};\n"); out.push_str("let final = new export:proxy {\n"); has_host = false; - for (name, link_type) in &self.exports.imports { + for (name, link_type) in &self.exports_links { match link_type { - LinkType::Debug => { - out.push_str(&format!("\"{name}\": debug[\"{name}\"] ,\n")); - } - LinkType::Recorder => { - if self.args.use_host_recorder { - has_host = true; - } else { - out.push_str(&format!("\"{name}\": recorder[\"{name}\"] ,\n")); - } - } - LinkType::Host => { - has_host = true; - } - LinkType::Imports => { - out.push_str(&format!("\"{name}\": imports[\"{name}\"] ,\n")); - } + LinkType::Imports => push_link(&mut out, name, "imports", name), LinkType::Main => { - if let Some(stripped) = name.strip_prefix("wrapped-") { - out.push_str(&format!("\"{name}\": main[\"{stripped}\"] ,\n")); - } else { - out.push_str(&format!("\"{name}\": main[\"{name}\"] ,\n")); - } + push_link(&mut out, name, "main", strip_wrapped(name).unwrap_or(name)) } + _ => has_host |= self.push_external_link(&mut out, name, link_type), } } - if has_host { - out.push_str("...,\n"); - } - out.push_str("};\n"); + close_instance(&mut out, has_host); out.push_str("export final...;\n"); std::fs::write(out_dir.join("compose.wac"), out.as_bytes())?; Ok(()) } + /// Links an import that is provided outside of the generated and main components. + /// Returns true if the import is left for the host to provide. + fn push_external_link(&self, out: &mut Source, name: &str, link_type: &LinkType) -> bool { + match link_type { + LinkType::Debug => push_link(out, name, "debug", name), + LinkType::Recorder if !self.args.use_host_recorder => { + push_link(out, name, "recorder", name) + } + LinkType::Recorder | LinkType::Host => return true, + LinkType::Imports | LinkType::Main => unreachable!(), + } + false + } pub fn generate_component( &mut self, resolve: &Resolve, id: WorldId, files: &mut Files, ) -> Result<()> { - self.generate_main_wit(resolve, id, files); + let main_pkg = resolve.worlds[id].package.unwrap(); + self.generate_main_wit(resolve, id, files)?; + self.generate_conversion_wit(resolve, main_pkg, files); + if matches!(self.mode, Mode::Record) { + generate_wrapped_wits(resolve, main_pkg, files)?; + } files.push( "deps/recorder.wit", include_str!("../assets/recorder.wit").as_bytes(), @@ -267,15 +215,13 @@ impl<'a> Opt<'a> { ); Ok(()) } - - pub fn generate_wrapped_wits(&self, dir: &std::path::Path) -> Result<()> { - let mut resolve = Resolve::default(); - let (main_id, _files) = resolve.push_dir(dir)?; - let has_version = package_with_version(&resolve); - // Generate conversion interface. Not updating resolve to avoid deep cloning the packages. + /// Generate the `proxy:conversion` interface for converting resources from + /// non-main packages between the generated components. + fn generate_conversion_wit(&self, resolve: &Resolve, main_pkg: PackageId, files: &mut Files) { + let has_version = package_with_version(resolve); let mut resources = BTreeMap::new(); for (_, iface) in resolve.interfaces.iter().filter(|(_, iface)| { - iface.package.is_some_and(|id| id != main_id) && iface.name.is_some() + iface.package.is_some_and(|id| id != main_pkg) && iface.name.is_some() }) { let pkg_id = iface.package.unwrap(); let pkg_name = &resolve.packages[pkg_id].name; @@ -308,17 +254,19 @@ impl<'a> Opt<'a> { let func_name = format!("{bindgen_name}-{resource}").to_kebab_case(); match self.mode { Mode::Record => { + let wrapped_iface = wrapped_name(&iface); + let wrapped_func = wrapped_name(&func_name); out.push_str(&format!( "\nuse {iface}.{{{resource} as host-{func_name}}};\n", )); out.push_str(&format!( - "use wrapped-{iface}.{{{resource} as wrapped-{func_name}}};\n", + "use {wrapped_iface}.{{{resource} as {wrapped_func}}};\n", )); out.push_str(&format!( - "get-wrapped-{func_name}: func(x: host-{func_name}) -> wrapped-{func_name};\n", + "get-{wrapped_func}: func(x: host-{func_name}) -> {wrapped_func};\n", )); out.push_str(&format!( - "get-host-{func_name}: func(x: wrapped-{func_name}) -> host-{func_name};\n", + "get-host-{func_name}: func(x: {wrapped_func}) -> host-{func_name};\n", )); } Mode::Replay | Mode::Fuzz | Mode::Dialog => { @@ -332,99 +280,110 @@ impl<'a> Opt<'a> { } } out.push_str("}\n"); - std::fs::write(dir.join("deps").join("conversion.wit"), out.as_bytes())?; - if matches!(self.mode, Mode::Record) { - // rename package name and generate wrapped wit - resolve.package_names = resolve - .package_names - .into_iter() - .map(|(mut name, id)| { - name.namespace = "wrapped-".to_string() + &name.namespace; - (name, id) - }) - .collect(); - for (_, pkg) in resolve.packages.iter_mut() { - pkg.name.namespace = "wrapped-".to_string() + &pkg.name.namespace; - } - for (id, pkg) in resolve.packages.iter().filter(|(id, _)| *id != main_id) { - let mut printer = WitPrinter::default(); - printer.print_package(&resolve, id, true)?; - let filename = if let Some(ver) = &pkg.name.version { - format!("wrapped-{}@{}.wit", pkg.name.name, ver) - } else { - format!("wrapped-{}.wit", pkg.name.name) - }; - std::fs::write(dir.join("deps").join(&filename), printer.output.to_string())?; - } - } - Ok(()) + files.push("deps/conversion.wit", out.as_bytes()); } fn load_imports(&mut self, file: &Path) -> Result<()> { - let (resolve, id) = load_wasm(file)?; - let world = &resolve.worlds[id]; - for (name, import) in &world.imports { - match import { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - let link_type = match name.as_str() { - "proxy:util/debug" => { - self.need_debug = true; - LinkType::Debug - } - "proxy:util/dialog" => LinkType::Host, - name if name.starts_with("proxy:recorder/") => LinkType::Recorder, - _ => LinkType::Host, - }; - self.imports.imports.insert(name.clone(), link_type); + for name in import_names(file)? { + let link_type = match name.as_str() { + "proxy:util/debug" => { + self.need_debug = true; + LinkType::Debug } - _ => todo!(), - } + "proxy:util/dialog" => LinkType::Host, + name if name.starts_with("proxy:recorder/") => LinkType::Recorder, + _ => LinkType::Host, + }; + self.imports_links.insert(name, link_type); } Ok(()) } fn load_exports(&mut self, file: &Path) -> Result<()> { - let (resolve, id) = load_wasm(file)?; - let world = &resolve.worlds[id]; - for (name, import) in &world.imports { - match import { - WorldItem::Interface { .. } => { - let name = resolve.name_world_key(name); - let link_type = match name.as_str() { - "proxy:util/debug" => { - self.need_debug = true; - LinkType::Debug - } - "proxy:conversion/conversion" => LinkType::Imports, - "proxy:util/dialog" => LinkType::Host, - name if name.starts_with("proxy:recorder/") => LinkType::Recorder, - name if matches!(self.mode, Mode::Record) => { - if let Some(stripped) = name.strip_prefix("wrapped-") { - if self.main.exports.contains(stripped) { - LinkType::Main - } else { - LinkType::Imports - } - } else { - LinkType::Host - } - } - name if self.main.exports.contains(name) => LinkType::Main, - _ => LinkType::Imports, - }; - self.exports.imports.insert(name.to_string(), link_type); + for name in import_names(file)? { + let link_type = match name.as_str() { + "proxy:util/debug" => { + self.need_debug = true; + LinkType::Debug } - _ => todo!(), - } + "proxy:conversion/conversion" => LinkType::Imports, + "proxy:util/dialog" => LinkType::Host, + name if name.starts_with("proxy:recorder/") => LinkType::Recorder, + name if matches!(self.mode, Mode::Record) => match strip_wrapped(name) { + Some(stripped) if self.main_exports.contains(stripped) => LinkType::Main, + Some(_) => LinkType::Imports, + None => LinkType::Host, + }, + name if self.main_exports.contains(name) => LinkType::Main, + _ => LinkType::Imports, + }; + self.exports_links.insert(name, link_type); } Ok(()) } } -fn load_wasm(file: &Path) -> Result<(Resolve, WorldId)> { +/// Generate a copy of every non-main package with the namespace prefixed by `wrapped-`. +fn generate_wrapped_wits(resolve: &Resolve, main_pkg: PackageId, files: &mut Files) -> Result<()> { + let mut resolve = resolve.clone(); + resolve.package_names = resolve + .package_names + .into_iter() + .map(|(mut name, id)| { + name.namespace = wrapped_name(&name.namespace); + (name, id) + }) + .collect(); + for (_, pkg) in resolve.packages.iter_mut() { + pkg.name.namespace = wrapped_name(&pkg.name.namespace); + } + for (id, pkg) in resolve.packages.iter().filter(|(id, _)| *id != main_pkg) { + let mut printer = WitPrinter::default(); + printer.print_package(&resolve, id, true)?; + let filename = match &pkg.name.version { + Some(ver) => format!("{}@{}.wit", pkg.name.name, ver), + None => format!("{}.wit", pkg.name.name), + }; + files.push( + &format!("deps/{}", wrapped_name(&filename)), + printer.output.to_string().as_bytes(), + ); + } + Ok(()) +} + +fn push_link(out: &mut Source, name: &str, instance: &str, export: &str) { + out.push_str(&format!("\"{name}\": {instance}[\"{export}\"] ,\n")); +} +fn close_instance(out: &mut Source, has_host: bool) { + if has_host { + out.push_str("...,\n"); + } + out.push_str("};\n"); +} + +/// Names of the interfaces in a world's imports or exports. +fn interface_names<'r>( + resolve: &Resolve, + items: impl IntoIterator, +) -> Result> { + items + .into_iter() + .map(|(key, item)| match item { + WorldItem::Interface { .. } => Ok(resolve.name_world_key(key)), + _ => bail!( + "unsupported world item `{}`: only interfaces are supported", + resolve.name_world_key(key) + ), + }) + .collect() +} + +/// Interface imports of a built component. Reading them from the wasm +/// accounts for unused imports being optimized away. +fn import_names(file: &Path) -> Result> { use wit_parser::decoding::{DecodedWasm, decode}; let bytes = std::fs::read(file)?; let DecodedWasm::Component(resolve, id) = decode(&bytes)? else { - panic!() + bail!("{} is not a component", file.display()); }; - Ok((resolve, id)) + interface_names(&resolve, &resolve.worlds[id].imports) } diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index 76e1822..d593ec3 100644 --- a/src/codegen/mod.rs +++ b/src/codegen/mod.rs @@ -1,4 +1,6 @@ -use crate::util::{FullTypePath, get_resource_from_trait_name, get_return_type, make_path}; +use crate::util::{ + FullTypePath, get_resource_from_trait_name, get_return_type, make_path, toggle_wrapped_module, +}; use anyhow::Result; use quote::quote; use std::collections::{BTreeMap, BTreeSet}; @@ -196,9 +198,7 @@ impl State { fn generate_conversion_func(&self, sig: &Signature) -> syn::ImplItemFn { let func_name = &sig.ident.to_string(); - let body = if func_name.starts_with("get_wrapped_") { - quote! { x.to_proxy() } - } else if func_name.starts_with("get_host_") { + let body = if func_name.starts_with("get_wrapped_") || func_name.starts_with("get_host_") { quote! { x.to_proxy() } } else if func_name.starts_with("get_mock_") { let resource = get_return_type(&sig.output).unwrap(); @@ -237,9 +237,6 @@ pub fn get_proxy_path(src_path: &[String]) -> Vec { res.insert(0, "exports".to_string()); wrapped_idx = 1; } - match res[wrapped_idx].strip_prefix("wrapped_") { - Some(name) => res[wrapped_idx] = name.to_string(), - None => res[wrapped_idx] = "wrapped_".to_string() + &res[wrapped_idx], - } + res[wrapped_idx] = toggle_wrapped_module(&res[wrapped_idx]); res } diff --git a/src/codegen/record.rs b/src/codegen/record.rs index 2671b3f..015359d 100644 --- a/src/codegen/record.rs +++ b/src/codegen/record.rs @@ -1,5 +1,7 @@ use super::{GenerateMode, State, get_proxy_path}; -use crate::util::{ResourceFuncKind, extract_arg_info, get_return_type, wit_func_name}; +use crate::util::{ + ResourceFuncKind, extract_arg_info, get_return_type, is_wrapped_module, wit_func_name, +}; use quote::quote; use syn::{Signature, parse_quote}; @@ -61,7 +63,7 @@ impl State { quote! { mut } }; let display_name = wit_func_name(module_path, resource, func_name, &kind); - let is_export = !module_path[1].starts_with("wrapped_"); + let is_export = !is_wrapped_module(&module_path[1]); let record_ret = if get_return_type(&sig.output).is_none() { quote! { #func(#(#call_args),*); diff --git a/src/instrument.rs b/src/instrument.rs index 8c2a550..2e736ea 100644 --- a/src/instrument.rs +++ b/src/instrument.rs @@ -44,7 +44,6 @@ pub fn run(args: InstrumentArgs) -> Result<()> { // 3. Parse the main wit file from tmp_dir/wit and feed into opts.generate_component let (resolve, world) = parse_wit(&wit_dir, None)?; let mut opts = crate::ast::Opt::new(&args); - opts.generate_wrapped_wits(&wit_dir)?; let mut files = Files::default(); opts.generate_component(&resolve, world, &mut files)?; @@ -55,7 +54,7 @@ pub fn run(args: InstrumentArgs) -> Result<()> { } // Re-generate exports world to bring in extra imports let (export_resolve, export_world) = parse_wit(&wit_dir, Some("tmp-exports"))?; - opts.generate_exports_world(&export_resolve, export_world, &mut files); + opts.generate_exports_world(&export_resolve, export_world, &mut files)?; for (name, content) in files.iter() { let path = wit_dir.as_path().join(name); eprintln!("Generating: {}", path.display()); diff --git a/src/traits/proxy.rs b/src/traits/proxy.rs index fd3d97a..ac05150 100644 --- a/src/traits/proxy.rs +++ b/src/traits/proxy.rs @@ -1,5 +1,6 @@ +use crate::codegen::State; use crate::traits::Trait; -use crate::{codegen::State, util::make_path}; +use crate::util::{is_wrapped_module, make_path}; use heck::ToSnakeCase; use quote::quote; use syn::{File, Item, ItemEnum, ItemStruct, parse_quote}; @@ -23,7 +24,7 @@ impl Trait for ProxyTrait<'_> { if in_import { let is_import_only = output_path[0] != "exports"; if is_import_only { - let call = if output_path[0].starts_with("wrapped_") { + let call = if is_wrapped_module(&output_path[0]) { "get_" } else { "get_host_" diff --git a/src/util.rs b/src/util.rs index cbe5938..44ead11 100644 --- a/src/util.rs +++ b/src/util.rs @@ -52,10 +52,7 @@ pub fn wit_func_name( Some(ResourceFuncKind::Method) => res.push_str("[method]"), _ => {} } - let nonwrapped = module_path[0] - .strip_prefix("wrapped_") - .unwrap_or(&module_path[0]); - res.push_str(&nonwrapped.to_kebab_case()); + res.push_str(&unwrapped_module(&module_path[0]).to_kebab_case()); res.push(':'); res.push_str(&module_path[1].to_kebab_case()); res.push('/'); @@ -378,6 +375,33 @@ fn extract_bitflag(macro_item: &syn::ItemMacro) -> Option { None } +// In record mode, the guest sees a copy of each imported package with its namespace +// prefixed by `wrapped-` in WIT, which becomes `wrapped_` in the generated Rust modules. +const WRAPPED_WIT: &str = "wrapped-"; +const WRAPPED_RUST: &str = "wrapped_"; +/// `ns:pkg/iface` -> `wrapped-ns:pkg/iface` +pub fn wrapped_name(name: &str) -> String { + format!("{WRAPPED_WIT}{name}") +} +/// `wrapped-ns:pkg/iface` -> `Some("ns:pkg/iface")` +pub fn strip_wrapped(name: &str) -> Option<&str> { + name.strip_prefix(WRAPPED_WIT) +} +pub fn is_wrapped_module(module: &str) -> bool { + module.starts_with(WRAPPED_RUST) +} +/// Strips the `wrapped_` prefix from a Rust module name if present. +pub fn unwrapped_module(module: &str) -> &str { + module.strip_prefix(WRAPPED_RUST).unwrap_or(module) +} +/// Maps a Rust module name between its host and wrapped copy. +pub fn toggle_wrapped_module(module: &str) -> String { + match module.strip_prefix(WRAPPED_RUST) { + Some(name) => name.to_string(), + None => format!("{WRAPPED_RUST}{module}"), + } +} + // utils for WIT names pub fn ident(name: &str) -> Cow<'_, str> { if is_keyword(name) { From 8ae581e65e8bf01bef0f718590b6873accf48041 Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 13:07:43 -0700 Subject: [PATCH 2/2] remove magic42 --- src/ast.rs | 4 +--- src/codegen/mod.rs | 36 +++++++++++++++++++++++++++++++----- src/traits/dialog.rs | 8 ++------ src/traits/fuzz.rs | 8 ++------ src/traits/wave.rs | 8 ++------ src/util.rs | 6 ++++++ 6 files changed, 44 insertions(+), 26 deletions(-) diff --git a/src/ast.rs b/src/ast.rs index d7bba58..81d220d 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -270,11 +270,9 @@ impl<'a> Opt<'a> { )); } Mode::Replay | Mode::Fuzz | Mode::Dialog => { - // Add a magic separator so that codegen::generate_conversion_func can recover the resource name - let magic_name = format!("{bindgen_name}-magic42-{resource}").to_kebab_case(); out.push_str(&format!("\nuse {iface}.{{{resource} as {func_name}}};\n")); out.push_str(&format!( - "get-mock-{magic_name}: func(handle: u32) -> {func_name};\n" + "get-mock-{func_name}: func(handle: u32) -> {func_name};\n" )); } } diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index d593ec3..80ef4ad 100644 --- a/src/codegen/mod.rs +++ b/src/codegen/mod.rs @@ -2,11 +2,12 @@ use crate::util::{ FullTypePath, get_resource_from_trait_name, get_return_type, make_path, toggle_wrapped_module, }; use anyhow::Result; +use heck::ToSnakeCase; use quote::quote; use std::collections::{BTreeMap, BTreeSet}; use std::path::PathBuf; use syn::{ - File, Ident, Item, ItemEnum, ItemStruct, ItemTrait, Signature, TraitItem, parse_quote, + File, Ident, Item, ItemEnum, ItemStruct, ItemTrait, Signature, TraitItem, Type, parse_quote, visit_mut::VisitMut, }; @@ -55,6 +56,8 @@ pub struct State { pub types: BTreeMap, Vec>, pub funcs: BTreeMap, BTreeMap, Vec>>, pub module_paths: BTreeSet>, + /// Full path of each public type alias to its target type. + pub type_aliases: BTreeMap, Type>, pub output: Vec, } pub enum TypeInfo { @@ -79,6 +82,7 @@ impl GenerateArgs { types: BTreeMap::new(), funcs: BTreeMap::new(), module_paths: BTreeSet::new(), + type_aliases: BTreeMap::new(), output: Vec::new(), }; state.generate_preamble(); @@ -202,10 +206,7 @@ impl State { quote! { x.to_proxy() } } else if func_name.starts_with("get_mock_") { let resource = get_return_type(&sig.output).unwrap(); - let name = func_name - .rfind("_magic42_") - .map(|idx| &func_name[idx + 9..]) - .unwrap(); + let name = self.resource_name(&resource); quote! { #resource::new(MockedResource { handle, name: #name.to_string() }) } @@ -218,6 +219,31 @@ impl State { } } } + /// The conversion interface refers to resources through type aliases, e.g. + /// `WasiIoStreamsOutputStream = ...::wasi::io::streams::OutputStream`. + /// Returns the snake case name of the aliased resource, e.g. `output_stream`. + fn resource_name(&self, alias: &Type) -> String { + let Type::Path(alias) = alias else { + unreachable!() + }; + let alias_path: Vec<_> = alias + .path + .segments + .iter() + .map(|s| s.ident.to_string()) + .collect(); + let Some(Type::Path(target)) = self.type_aliases.get(&alias_path) else { + panic!("{} is not a type alias", alias_path.join("::")); + }; + target + .path + .segments + .last() + .unwrap() + .ident + .to_string() + .to_snake_case() + } fn into_output_file(self) -> File { File { items: self.output, diff --git a/src/traits/dialog.rs b/src/traits/dialog.rs index 44b7127..8912e96 100644 --- a/src/traits/dialog.rs +++ b/src/traits/dialog.rs @@ -13,12 +13,8 @@ impl Trait for DialogTrait { let wit_name = resource.ident.to_string().to_kebab_case(); let in_import = module_path[0] != "exports"; if in_import { - let call = format!( - "get_mock_{}_magic42_{}", - module_path.join("_"), - resource.ident - ) - .to_snake_case(); + let call = + format!("get_mock_{}_{}", module_path.join("_"), resource.ident).to_snake_case(); let call: syn::Ident = syn::parse_str(&call).unwrap(); res.push(parse_quote! { impl Dialog for #resource_path { diff --git a/src/traits/fuzz.rs b/src/traits/fuzz.rs index 5d3ff2a..b20d16e 100644 --- a/src/traits/fuzz.rs +++ b/src/traits/fuzz.rs @@ -13,12 +13,8 @@ impl Trait for FuzzTrait { let wit_name = resource.ident.to_string().to_kebab_case(); let in_import = module_path[0] != "exports"; if in_import { - let call = format!( - "get_mock_{}_magic42_{}", - module_path.join("_"), - resource.ident - ) - .to_snake_case(); + let call = + format!("get_mock_{}_{}", module_path.join("_"), resource.ident).to_snake_case(); let call: syn::Ident = syn::parse_str(&call).unwrap(); res.push(parse_quote! { impl Arbitrary<'_> for #resource_path { diff --git a/src/traits/wave.rs b/src/traits/wave.rs index 6eb3210..1c6c9dc 100644 --- a/src/traits/wave.rs +++ b/src/traits/wave.rs @@ -50,12 +50,8 @@ impl Trait for WaveTrait { }); } if self.to_rust { - let call = format!( - "get_mock_{}_magic42_{}", - module_path.join("_"), - resource.ident - ) - .to_snake_case(); + let call = format!("get_mock_{}_{}", module_path.join("_"), resource.ident) + .to_snake_case(); let call: syn::Ident = syn::parse_str(&call).unwrap(); res.push(parse_quote! { impl ToRust<#resource_path> for Value { diff --git a/src/util.rs b/src/util.rs index 44ead11..b4bdd9d 100644 --- a/src/util.rs +++ b/src/util.rs @@ -235,6 +235,12 @@ impl codegen::State { .or_default() .push(TypeInfo::Enum(enum_item)); } + // type aliases, e.g. the resources used by the conversion interface + Item::Type(type_item) if matches!(type_item.vis, Visibility::Public(_)) => { + let mut path = current_path.clone(); + path.push(type_item.ident.to_string()); + self.type_aliases.insert(path, (*type_item.ty).clone()); + } // flags Item::Macro(macro_item) => { if let Some(enum_item) = extract_bitflag(macro_item) {