refactor

src/codegen/mod.rs

@@ -1,10 +1,13 @@
-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 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,
 };
 
@@ -53,6 +56,8 @@ pub struct State {
     pub types: BTreeMap<Vec<String>, Vec<TypeInfo>>,
     pub funcs: BTreeMap<Vec<String>, BTreeMap<Option<String>, Vec<Signature>>>,
     pub module_paths: BTreeSet<Vec<String>>,
+    /// Full path of each public type alias to its target type.
+    pub type_aliases: BTreeMap<Vec<String>, Type>,
     pub output: Vec<Item>,
 }
 pub enum TypeInfo {
@@ -77,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();
@@ -196,16 +202,11 @@ 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();
-            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,
@@ -237,9 +263,6 @@ pub fn get_proxy_path(src_path: &[String]) -> Vec<String> {
         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
 }

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),*);

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());

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 {

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 {

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_"

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 {

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('/');
@@ -238,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) {
@@ -378,6 +381,33 @@ fn extract_bitflag(macro_item: &syn::ItemMacro) -> Option<ItemFlag> {
     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) {