From 3bc402cb83fe0919a206577cdd9c13afcb5b87ed Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 18:55:28 -0700 Subject: [PATCH 1/4] split out util.rs --- src/ast.rs | 2 +- src/codegen/dialog.rs | 2 +- src/codegen/fuzz.rs | 2 +- src/codegen/mod.rs | 8 +- src/codegen/parse.rs | 209 +++++++++++++++ src/codegen/record.rs | 4 +- src/codegen/replay.rs | 2 +- src/codegen/util.rs | 194 ++++++++++++++ src/instrument.rs | 4 +- src/main.rs | 2 +- src/traits/dialog.rs | 2 +- src/traits/fuzz.rs | 2 +- src/traits/proxy.rs | 2 +- src/traits/wave.rs | 2 +- src/traits/wit.rs | 2 +- src/util.rs | 580 ------------------------------------------ src/wit_util.rs | 178 +++++++++++++ 17 files changed, 600 insertions(+), 597 deletions(-) create mode 100644 src/codegen/parse.rs create mode 100644 src/codegen/util.rs delete mode 100644 src/util.rs create mode 100644 src/wit_util.rs diff --git a/src/ast.rs b/src/ast.rs index 81d220d..d3ba2b3 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -1,6 +1,6 @@ use crate::Mode; use crate::instrument::InstrumentArgs; -use crate::util::*; +use crate::wit_util::*; use anyhow::{Result, bail}; use std::collections::{BTreeMap, BTreeSet}; use std::path::Path; diff --git a/src/codegen/dialog.rs b/src/codegen/dialog.rs index c3c7ea5..ee0dfe3 100644 --- a/src/codegen/dialog.rs +++ b/src/codegen/dialog.rs @@ -1,5 +1,5 @@ use super::State; -use crate::util::{ +use super::util::{ FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, get_return_type, make_path, wit_func_name, }; diff --git a/src/codegen/fuzz.rs b/src/codegen/fuzz.rs index 28a2d5e..33849b9 100644 --- a/src/codegen/fuzz.rs +++ b/src/codegen/fuzz.rs @@ -1,5 +1,5 @@ use super::State; -use crate::util::{ +use super::util::{ FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, get_return_type, make_path, wit_func_name, }; diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index 0ead336..caecd8b 100644 --- a/src/codegen/mod.rs +++ b/src/codegen/mod.rs @@ -1,6 +1,3 @@ -use crate::util::{ - FullTypePath, get_resource_from_trait_name, get_return_type, make_path, toggle_wrapped_module, -}; use anyhow::Result; use heck::ToKebabCase; use quote::quote; @@ -10,11 +7,16 @@ use syn::{ File, Ident, Item, ItemEnum, ItemStruct, ItemTrait, Signature, TraitItem, Type, parse_quote, visit_mut::VisitMut, }; +use util::{ + FullTypePath, get_resource_from_trait_name, get_return_type, make_path, toggle_wrapped_module, +}; mod dialog; mod fuzz; +mod parse; mod record; mod replay; +pub mod util; #[derive(clap::Parser)] pub struct GenerateArgs { diff --git a/src/codegen/parse.rs b/src/codegen/parse.rs new file mode 100644 index 0000000..03668f7 --- /dev/null +++ b/src/codegen/parse.rs @@ -0,0 +1,209 @@ +use super::util::{FullTypePath, get_resource_from_trait_name}; +use super::{ItemFlag, State, TypeInfo}; +use quote::ToTokens; +use syn::{Ident, Item, Signature, Type, Visibility, visit_mut::VisitMut}; + +impl State { + pub fn find_all_items(&mut self, items: &[Item], current_path: Vec) { + for item in items { + match item { + // export functions + Item::Trait(trait_item) if matches!(trait_item.vis, Visibility::Public(_)) => { + let trait_item = trait_item.clone(); + let resource = get_resource_from_trait_name(&trait_item.ident.to_string()); + let funcs = self.funcs_mut(¤t_path, resource); + for item in &trait_item.items { + if let syn::TraitItem::Fn(method) = item { + funcs.push(method.sig.clone()); + } + } + self.traits + .entry(current_path.clone()) + .or_default() + .push(trait_item); + } + // import resource functions + Item::Impl(impl_item) if impl_item.trait_.is_none() => { + let resource = if let Type::Path(type_path) = &*impl_item.self_ty { + type_path.path.segments.last().unwrap().ident.to_string() + } else { + unreachable!() + }; + let funcs = self.funcs_mut(¤t_path, Some(resource)); + for item in &impl_item.items { + if let syn::ImplItem::Fn(method) = item { + if !matches!(method.vis, Visibility::Public(_)) { + continue; + } + if has_doc_hidden(&method.attrs) { + continue; + } + funcs.push(method.sig.clone()); + } + } + } + // import top-level functions + Item::Fn(func_item) + if matches!(func_item.vis, Visibility::Public(_)) + && current_path.len() >= 3 => + { + self.funcs_mut(¤t_path, None) + .push(func_item.sig.clone()); + } + // resource and struct types + Item::Struct(struct_item) if matches!(struct_item.vis, Visibility::Public(_)) => { + let has_repr_transparent = struct_item.attrs.iter().any(|attr| { + attr.path().is_ident("repr") + && attr.to_token_stream().to_string().contains("transparent") + }); + let type_info = if has_repr_transparent { + if struct_item.ident.to_string().ends_with("Borrow") + && current_path[0] == "exports" + { + continue; + } + TypeInfo::Resource(struct_item.clone()) + } else { + let mut struct_item = struct_item.clone(); + let mut transformer = FullTypePath { + module_path: ¤t_path, + }; + transformer.visit_item_struct_mut(&mut struct_item); + TypeInfo::Struct(struct_item) + }; + self.push_type(¤t_path, type_info); + } + // enum types + Item::Enum(enum_item) if matches!(enum_item.vis, Visibility::Public(_)) => { + let mut enum_item = enum_item.clone(); + let mut transformer = FullTypePath { + module_path: ¤t_path, + }; + transformer.visit_item_enum_mut(&mut enum_item); + self.push_type(¤t_path, 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) { + self.push_type(¤t_path, TypeInfo::Flag(enum_item)); + } + } + // traverse down the modules + Item::Mod(module) if matches!(module.vis, Visibility::Public(_)) => { + if let Some((_, ref mod_items)) = module.content { + let mut new_path = current_path.clone(); + let mod_name = module.ident.to_string(); + if current_path.is_empty() && mod_name == "_rt" { + continue; + } + new_path.push(mod_name); + self.module_paths.insert(new_path.clone()); + self.find_all_items(mod_items, new_path); + } + } + _ => {} + } + } + } + fn funcs_mut( + &mut self, + module_path: &[String], + resource: Option, + ) -> &mut Vec { + self.funcs + .entry(module_path.to_vec()) + .or_default() + .entry(resource) + .or_default() + } + fn push_type(&mut self, module_path: &[String], type_info: TypeInfo) { + self.types + .entry(module_path.to_vec()) + .or_default() + .push(type_info); + } + pub fn find_function( + &self, + module_path: &[String], + resource: &Option, + func: &Ident, + ) -> Option<&Signature> { + let module = self.funcs.get(module_path)?; + let funcs = module.get(resource)?; + funcs.iter().find(|sig| sig.ident == *func) + } + pub fn has_type_def(&self, module_path: &[String], name: &str) -> bool { + let types = match self.types.get(module_path) { + Some(types) => types, + None => return false, + }; + for type_info in types { + match type_info { + TypeInfo::Resource(struct_item) | TypeInfo::Struct(struct_item) => { + if struct_item.ident == name { + return true; + } + } + TypeInfo::Enum(enum_item) => { + if enum_item.ident == name { + return true; + } + } + TypeInfo::Flag(item_flag) => { + if item_flag.name == name { + return true; + } + } + } + } + false + } +} + +fn has_doc_hidden(attrs: &[syn::Attribute]) -> bool { + for attr in attrs { + if attr.path().is_ident("doc") + && let syn::Meta::List(meta_list) = &attr.meta + && meta_list.to_token_stream().to_string().contains("hidden") + { + return true; + } + } + false +} +fn extract_bitflag(macro_item: &syn::ItemMacro) -> Option { + use syn::{Attribute, Token, parse::Parser}; + if macro_item.mac.path.segments.last()?.ident == "bitflags" { + let tokens = macro_item.mac.tokens.clone(); + let parser = |input: syn::parse::ParseStream| { + input.call(Attribute::parse_outer)?; + input.parse::()?; + input.parse::()?; + let ident: Ident = input.parse()?; + input.parse::()?; + input.parse::()?; + let content; + syn::braced!(content in input); + let mut flags = vec![]; + while !content.is_empty() { + content.parse::()?; + let flag_ident: Ident = content.parse()?; + flags.push(flag_ident); + while !content.peek(Token![;]) && !content.is_empty() { + content.parse::()?; + } + content.parse::()?; + } + Ok((ident, flags)) + }; + let (name, flags) = parser.parse2(tokens).unwrap(); + return Some(ItemFlag { name, flags }); + } + None +} diff --git a/src/codegen/record.rs b/src/codegen/record.rs index 015359d..5bdbfaf 100644 --- a/src/codegen/record.rs +++ b/src/codegen/record.rs @@ -1,7 +1,7 @@ -use super::{GenerateMode, State, get_proxy_path}; -use crate::util::{ +use super::util::{ ResourceFuncKind, extract_arg_info, get_return_type, is_wrapped_module, wit_func_name, }; +use super::{GenerateMode, State, get_proxy_path}; use quote::quote; use syn::{Signature, parse_quote}; diff --git a/src/codegen/replay.rs b/src/codegen/replay.rs index 9212c9b..ae59f24 100644 --- a/src/codegen/replay.rs +++ b/src/codegen/replay.rs @@ -1,5 +1,5 @@ use super::State; -use crate::util::{ +use super::util::{ FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, get_return_type, make_path, wit_func_name, }; diff --git a/src/codegen/util.rs b/src/codegen/util.rs new file mode 100644 index 0000000..280be9f --- /dev/null +++ b/src/codegen/util.rs @@ -0,0 +1,194 @@ +use heck::ToKebabCase; +use syn::{FnArg, Ident, Signature, Type, parse_quote, visit_mut::VisitMut}; +pub struct FullTypePath<'a> { + pub module_path: &'a [String], +} +pub struct ArgInfo { + pub ident: Ident, + pub is_borrowed: bool, + pub ty: Type, +} +pub enum ResourceFuncKind { + Method, + Constructor, +} +impl ArgInfo { + pub fn call_param(&self) -> syn::Expr { + let ident = &self.ident; + if self.is_borrowed { + parse_quote! { &#ident } + } else { + parse_quote! { #ident } + } + } +} +pub fn make_path(module_path: &[String], name: &str) -> syn::Path { + let path = format!("{}::{}", module_path.join("::"), name); + syn::parse_str(&path).unwrap() +} +pub fn wit_func_name( + module_path: &[String], + resource: &Option, + func_name: &Ident, + kind: &Option, +) -> String { + assert!(module_path.len() >= 3); + let mut module_path = module_path.to_vec(); + if module_path[0] == "exports" { + module_path.remove(0); + } + assert!(module_path.len() == 3); + let mut res = String::new(); + match kind { + Some(ResourceFuncKind::Constructor) => res.push_str("[constructor]"), + Some(ResourceFuncKind::Method) => res.push_str("[method]"), + _ => {} + } + res.push_str(&unwrapped_module(&module_path[0]).to_kebab_case()); + res.push(':'); + res.push_str(&module_path[1].to_kebab_case()); + res.push('/'); + res.push_str(&module_path[2].to_kebab_case()); + if let Some(name) = resource { + res.push_str(&format!("/{}", name.to_kebab_case())); + } + res.push('.'); + res.push_str(&func_name.to_string().to_kebab_case()); + res +} +/// In virtualized components, constructors return `Self`, i.e. `MockedResource`, so type-directed +/// traits like `Dialog` or `Arbitrary` cannot recover the resource name. Returns the WIT name of +/// the resource to patch into the `MockedResource`, if `kind` is a constructor. +pub fn constructor_resource_name( + resource: &Option, + kind: &Option, +) -> Option { + match kind { + Some(ResourceFuncKind::Constructor) => Some(resource.as_ref()?.to_kebab_case()), + _ => None, + } +} +pub fn get_return_type(ret: &syn::ReturnType) -> Option { + match ret { + syn::ReturnType::Default => None, + syn::ReturnType::Type(_, ty) => match **ty { + Type::Tuple(ref tuple) if tuple.elems.is_empty() => None, + _ => Some(*ty.clone()), + }, + } +} + +pub fn extract_arg_info(sig: &Signature) -> (Option, Vec) { + let mut kind = None; + let mut arg_infos = Vec::new(); + for arg in sig.inputs.iter() { + match arg { + FnArg::Receiver(_) => { + kind = Some(ResourceFuncKind::Method); + } + FnArg::Typed(pat_type) => { + let ident = match &*pat_type.pat { + syn::Pat::Ident(ident) => ident.ident.clone(), + _ => unreachable!(), + }; + let ty = *pat_type.ty.clone(); + let is_borrowed = matches!(&*pat_type.ty, Type::Reference(_)); + arg_infos.push(ArgInfo { + ident, + is_borrowed, + ty, + }); + } + } + } + if sig.ident == "new" + && let Some(Type::Path(path)) = get_return_type(&sig.output) + && path.path.is_ident("Self") + { + kind = Some(ResourceFuncKind::Constructor); + } + (kind, arg_infos) +} + +const BUILTIN_TYPES: &[&str] = &[ + "Self", "Result", "Option", "Vec", "Box", "Rc", "Arc", "String", "str", "u8", "u16", "u32", + "u64", "u128", "usize", "i8", "i16", "i32", "i64", "i128", "isize", "f32", "f64", "bool", + "char", "_rt", +]; +impl<'a> VisitMut for FullTypePath<'a> { + fn visit_type_path_mut(&mut self, ty: &mut syn::TypePath) { + if ty.qself.is_none() && !ty.path.segments.is_empty() && ty.path.leading_colon.is_none() { + let ident = &ty.path.segments[0].ident.to_string(); + if BUILTIN_TYPES.contains(&ident.as_str()) { + if ident == "_rt" { + assert!(ty.path.segments.len() == 2); + ty.path.segments = ty.path.segments.iter().skip(1).cloned().collect(); + } + syn::visit_mut::visit_type_path_mut(self, ty); + return; + } + let module_idents = self + .module_path + .iter() + .map(|s| syn::parse_str::(s).unwrap()); + let original = &ty.path; + *ty = parse_quote! { + #(#module_idents)::*::#original + }; + } + syn::visit_mut::visit_type_path_mut(self, ty); + } +} + +pub fn get_resource_from_trait_name(trait_name: &str) -> Option { + let resource = trait_name.strip_prefix("Guest").unwrap(); + match resource { + "" => None, + name => Some(name.to_string()), + } +} + +pub fn get_owned_type(ty: &Type) -> Option { + match ty { + Type::Reference(type_ref) => { + match &*type_ref.elem { + Type::Slice(type_slice) => { + let inner_ty = &*type_slice.elem; + Some(parse_quote! { Vec<#inner_ty> }) + } + Type::Path(type_path) => { + if type_path.qself.is_none() + && type_path.path.segments.len() == 1 + && type_path.path.segments[0].ident == "str" + { + Some(parse_quote! { String }) + } else { + // TODO: need to handle nested borrow + Some(parse_quote! { #type_path }) + } + } + _ => None, + } + } + _ => None, + } +} + +// In record mode, the guest sees a copy of each imported package with its namespace +// prefixed by `wrapped-` in WIT (see `crate::wit_util::WRAPPED_WIT`), which becomes +// `wrapped_` in the generated Rust modules. +const WRAPPED_RUST: &str = "wrapped_"; +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}"), + } +} diff --git a/src/instrument.rs b/src/instrument.rs index 2e736ea..17e7207 100644 --- a/src/instrument.rs +++ b/src/instrument.rs @@ -39,7 +39,7 @@ pub fn run(args: InstrumentArgs) -> Result<()> { // 2. Extract WIT from the wasm component into {tmp_dir/wit} // TODO: no need to write to disk - crate::util::extract_wit(&args.wasm_file, &wit_dir)?; + crate::wit_util::extract_wit(&args.wasm_file, &wit_dir)?; // 3. Parse the main wit file from tmp_dir/wit and feed into opts.generate_component let (resolve, world) = parse_wit(&wit_dir, None)?; @@ -152,7 +152,7 @@ fn bindgen( use_custom_allocator: bool, ) -> Result<()> { let out_dir = tmp_dir.join(dest_name); - crate::util::generate_bindings(wit_dir, world_name, &out_dir)?; + crate::wit_util::generate_bindings(wit_dir, world_name, &out_dir)?; let binding_file = out_dir.join(world_name.to_owned() + ".rs"); let codegen_mode = match mode { Mode::Record => codegen::GenerateMode::Record, diff --git a/src/main.rs b/src/main.rs index f0748eb..4dc4273 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,7 @@ mod ast; mod codegen; mod instrument; mod traits; -mod util; +mod wit_util; #[cfg(feature = "run")] mod run; diff --git a/src/traits/dialog.rs b/src/traits/dialog.rs index 8912e96..6847b2b 100644 --- a/src/traits/dialog.rs +++ b/src/traits/dialog.rs @@ -1,5 +1,5 @@ +use crate::codegen::util::{make_path, wit_func_name}; use crate::traits::Trait; -use crate::util::{make_path, wit_func_name}; use heck::{ToKebabCase, ToSnakeCase}; use quote::quote; use syn::{Item, ItemEnum, ItemStruct, parse_quote}; diff --git a/src/traits/fuzz.rs b/src/traits/fuzz.rs index b20d16e..cb62604 100644 --- a/src/traits/fuzz.rs +++ b/src/traits/fuzz.rs @@ -1,5 +1,5 @@ +use crate::codegen::util::make_path; use crate::traits::Trait; -use crate::util::make_path; use heck::{ToKebabCase, ToSnakeCase}; use quote::quote; use syn::{Item, ItemEnum, ItemStruct, parse_quote}; diff --git a/src/traits/proxy.rs b/src/traits/proxy.rs index f6995d7..9634ddd 100644 --- a/src/traits/proxy.rs +++ b/src/traits/proxy.rs @@ -1,6 +1,6 @@ use crate::codegen::State; +use crate::codegen::util::{is_wrapped_module, make_path}; use crate::traits::Trait; -use crate::util::{is_wrapped_module, make_path}; use heck::ToSnakeCase; use proc_macro2::TokenStream; use quote::{format_ident, quote}; diff --git a/src/traits/wave.rs b/src/traits/wave.rs index 1c6c9dc..519613c 100644 --- a/src/traits/wave.rs +++ b/src/traits/wave.rs @@ -1,5 +1,5 @@ +use crate::codegen::util::make_path; use crate::traits::Trait; -use crate::util::make_path; use heck::{ToKebabCase, ToSnakeCase}; use quote::quote; use syn::{Item, ItemEnum, ItemStruct, parse_quote}; diff --git a/src/traits/wit.rs b/src/traits/wit.rs index 002e85b..317dc62 100644 --- a/src/traits/wit.rs +++ b/src/traits/wit.rs @@ -1,5 +1,5 @@ +use crate::codegen::util::{make_path, wit_func_name}; use crate::traits::Trait; -use crate::util::{make_path, wit_func_name}; use heck::ToKebabCase; use syn::{Item, ItemEnum, ItemStruct, parse_quote}; diff --git a/src/util.rs b/src/util.rs deleted file mode 100644 index 2219ee4..0000000 --- a/src/util.rs +++ /dev/null @@ -1,580 +0,0 @@ -use crate::codegen::{self, ItemFlag, TypeInfo}; -use anyhow::{Context, Result}; -use heck::ToKebabCase; -use quote::ToTokens; -use std::borrow::Cow; -use std::collections::{BTreeMap, BTreeSet, HashMap}; -use std::path::Path; -use syn::{FnArg, Ident, Item, Signature, Type, Visibility, parse_quote, visit_mut::VisitMut}; -use wit_component::WitPrinter; -use wit_parser::{PackageId, Resolve}; -pub struct FullTypePath<'a> { - pub module_path: &'a [String], -} -pub struct ArgInfo { - pub ident: Ident, - pub is_borrowed: bool, - pub ty: Type, -} -pub enum ResourceFuncKind { - Method, - Constructor, -} -impl ArgInfo { - pub fn call_param(&self) -> syn::Expr { - let ident = &self.ident; - if self.is_borrowed { - parse_quote! { &#ident } - } else { - parse_quote! { #ident } - } - } -} -pub fn make_path(module_path: &[String], name: &str) -> syn::Path { - let path = format!("{}::{}", module_path.join("::"), name); - syn::parse_str(&path).unwrap() -} -pub fn wit_func_name( - module_path: &[String], - resource: &Option, - func_name: &Ident, - kind: &Option, -) -> String { - assert!(module_path.len() >= 3); - let mut module_path = module_path.to_vec(); - if module_path[0] == "exports" { - module_path.remove(0); - } - assert!(module_path.len() == 3); - let mut res = String::new(); - match kind { - Some(ResourceFuncKind::Constructor) => res.push_str("[constructor]"), - Some(ResourceFuncKind::Method) => res.push_str("[method]"), - _ => {} - } - res.push_str(&unwrapped_module(&module_path[0]).to_kebab_case()); - res.push(':'); - res.push_str(&module_path[1].to_kebab_case()); - res.push('/'); - res.push_str(&module_path[2].to_kebab_case()); - if let Some(name) = resource { - res.push_str(&format!("/{}", name.to_kebab_case())); - } - res.push('.'); - res.push_str(&func_name.to_string().to_kebab_case()); - res -} -/// In virtualized components, constructors return `Self`, i.e. `MockedResource`, so type-directed -/// traits like `Dialog` or `Arbitrary` cannot recover the resource name. Returns the WIT name of -/// the resource to patch into the `MockedResource`, if `kind` is a constructor. -pub fn constructor_resource_name( - resource: &Option, - kind: &Option, -) -> Option { - match kind { - Some(ResourceFuncKind::Constructor) => Some(resource.as_ref()?.to_kebab_case()), - _ => None, - } -} -pub fn get_return_type(ret: &syn::ReturnType) -> Option { - match ret { - syn::ReturnType::Default => None, - syn::ReturnType::Type(_, ty) => match **ty { - Type::Tuple(ref tuple) if tuple.elems.is_empty() => None, - _ => Some(*ty.clone()), - }, - } -} - -pub fn extract_arg_info(sig: &Signature) -> (Option, Vec) { - let mut kind = None; - let mut arg_infos = Vec::new(); - for arg in sig.inputs.iter() { - match arg { - FnArg::Receiver(_) => { - kind = Some(ResourceFuncKind::Method); - } - FnArg::Typed(pat_type) => { - let ident = match &*pat_type.pat { - syn::Pat::Ident(ident) => ident.ident.clone(), - _ => unreachable!(), - }; - let ty = *pat_type.ty.clone(); - let is_borrowed = matches!(&*pat_type.ty, Type::Reference(_)); - arg_infos.push(ArgInfo { - ident, - is_borrowed, - ty, - }); - } - } - } - if sig.ident == "new" - && let Some(Type::Path(path)) = get_return_type(&sig.output) - && path.path.is_ident("Self") - { - kind = Some(ResourceFuncKind::Constructor); - } - (kind, arg_infos) -} - -const BUILTIN_TYPES: &[&str] = &[ - "Self", "Result", "Option", "Vec", "Box", "Rc", "Arc", "String", "str", "u8", "u16", "u32", - "u64", "u128", "usize", "i8", "i16", "i32", "i64", "i128", "isize", "f32", "f64", "bool", - "char", "_rt", -]; -impl<'a> VisitMut for FullTypePath<'a> { - fn visit_type_path_mut(&mut self, ty: &mut syn::TypePath) { - if ty.qself.is_none() && !ty.path.segments.is_empty() && ty.path.leading_colon.is_none() { - let ident = &ty.path.segments[0].ident.to_string(); - if BUILTIN_TYPES.contains(&ident.as_str()) { - if ident == "_rt" { - assert!(ty.path.segments.len() == 2); - ty.path.segments = ty.path.segments.iter().skip(1).cloned().collect(); - } - syn::visit_mut::visit_type_path_mut(self, ty); - return; - } - let module_idents = self - .module_path - .iter() - .map(|s| syn::parse_str::(s).unwrap()); - let original = &ty.path; - *ty = parse_quote! { - #(#module_idents)::*::#original - }; - } - syn::visit_mut::visit_type_path_mut(self, ty); - } -} - -impl codegen::State { - pub fn find_all_items(&mut self, items: &[Item], current_path: Vec) { - for item in items { - match item { - // export functions - Item::Trait(trait_item) if matches!(trait_item.vis, Visibility::Public(_)) => { - let trait_item = trait_item.clone(); - let resource = get_resource_from_trait_name(&trait_item.ident.to_string()); - let funcs = self - .funcs - .entry(current_path.clone()) - .or_default() - .entry(resource) - .or_default(); - for item in &trait_item.items { - if let syn::TraitItem::Fn(method) = item { - funcs.push(method.sig.clone()); - } - } - self.traits - .entry(current_path.clone()) - .or_default() - .push(trait_item); - } - // import resource functions - Item::Impl(impl_item) if impl_item.trait_.is_none() => { - let resource = if let Type::Path(type_path) = &*impl_item.self_ty { - type_path.path.segments.last().unwrap().ident.to_string() - } else { - unreachable!() - }; - let funcs = self - .funcs - .entry(current_path.clone()) - .or_default() - .entry(Some(resource)) - .or_default(); - for item in &impl_item.items { - if let syn::ImplItem::Fn(method) = item { - if !matches!(method.vis, Visibility::Public(_)) { - continue; - } - if has_doc_hidden(&method.attrs) { - continue; - } - funcs.push(method.sig.clone()); - } - } - } - // import top-level functions - Item::Fn(func_item) - if matches!(func_item.vis, Visibility::Public(_)) - && current_path.len() >= 3 => - { - self.funcs - .entry(current_path.clone()) - .or_default() - .entry(None) - .or_default() - .push(func_item.sig.clone()); - } - // resource and struct types - Item::Struct(struct_item) if matches!(struct_item.vis, Visibility::Public(_)) => { - let has_repr_transparent = struct_item.attrs.iter().any(|attr| { - attr.path().is_ident("repr") - && attr.to_token_stream().to_string().contains("transparent") - }); - let type_info = if has_repr_transparent { - if struct_item.ident.to_string().ends_with("Borrow") - && current_path[0] == "exports" - { - continue; - } - TypeInfo::Resource(struct_item.clone()) - } else { - let mut struct_item = struct_item.clone(); - let mut transformer = FullTypePath { - module_path: ¤t_path, - }; - transformer.visit_item_struct_mut(&mut struct_item); - TypeInfo::Struct(struct_item) - }; - self.types - .entry(current_path.clone()) - .or_default() - .push(type_info); - } - // enum types - Item::Enum(enum_item) if matches!(enum_item.vis, Visibility::Public(_)) => { - let mut enum_item = enum_item.clone(); - let mut transformer = FullTypePath { - module_path: ¤t_path, - }; - transformer.visit_item_enum_mut(&mut enum_item); - self.types - .entry(current_path.clone()) - .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) { - self.types - .entry(current_path.clone()) - .or_default() - .push(TypeInfo::Flag(enum_item)); - } - } - // traverse down the modules - Item::Mod(module) if matches!(module.vis, Visibility::Public(_)) => { - if let Some((_, ref mod_items)) = module.content { - let mut new_path = current_path.clone(); - let mod_name = module.ident.to_string(); - if current_path.is_empty() && mod_name == "_rt" { - continue; - } - new_path.push(mod_name); - self.module_paths.insert(new_path.clone()); - self.find_all_items(mod_items, new_path); - } - } - _ => {} - } - } - } - pub fn find_function( - &self, - module_path: &[String], - resource: &Option, - func: &Ident, - ) -> Option<&Signature> { - let module = self.funcs.get(module_path)?; - let funcs = module.get(resource)?; - funcs.iter().find(|sig| sig.ident == *func) - } - pub fn has_type_def(&self, module_path: &[String], name: &str) -> bool { - let types = match self.types.get(module_path) { - Some(types) => types, - None => return false, - }; - for type_info in types { - match type_info { - TypeInfo::Resource(struct_item) | TypeInfo::Struct(struct_item) => { - if struct_item.ident == name { - return true; - } - } - TypeInfo::Enum(enum_item) => { - if enum_item.ident == name { - return true; - } - } - TypeInfo::Flag(item_flag) => { - if item_flag.name == name { - return true; - } - } - } - } - false - } -} - -pub fn get_resource_from_trait_name(trait_name: &str) -> Option { - let resource = trait_name.strip_prefix("Guest").unwrap(); - match resource { - "" => None, - name => Some(name.to_string()), - } -} - -pub fn get_owned_type(ty: &Type) -> Option { - match ty { - Type::Reference(type_ref) => { - match &*type_ref.elem { - Type::Slice(type_slice) => { - let inner_ty = &*type_slice.elem; - Some(parse_quote! { Vec<#inner_ty> }) - } - Type::Path(type_path) => { - if type_path.qself.is_none() - && type_path.path.segments.len() == 1 - && type_path.path.segments[0].ident == "str" - { - Some(parse_quote! { String }) - } else { - // TODO: need to handle nested borrow - Some(parse_quote! { #type_path }) - } - } - _ => None, - } - } - _ => None, - } -} - -fn has_doc_hidden(attrs: &[syn::Attribute]) -> bool { - for attr in attrs { - if attr.path().is_ident("doc") - && let syn::Meta::List(meta_list) = &attr.meta - && meta_list.to_token_stream().to_string().contains("hidden") - { - return true; - } - } - false -} -fn extract_bitflag(macro_item: &syn::ItemMacro) -> Option { - use syn::{Attribute, Token, parse::Parser}; - if macro_item.mac.path.segments.last()?.ident == "bitflags" { - let tokens = macro_item.mac.tokens.clone(); - let parser = |input: syn::parse::ParseStream| { - input.call(Attribute::parse_outer)?; - input.parse::()?; - input.parse::()?; - let ident: Ident = input.parse()?; - input.parse::()?; - input.parse::()?; - let content; - syn::braced!(content in input); - let mut flags = vec![]; - while !content.is_empty() { - content.parse::()?; - let flag_ident: Ident = content.parse()?; - flags.push(flag_ident); - while !content.peek(Token![;]) && !content.is_empty() { - content.parse::()?; - } - content.parse::()?; - } - Ok((ident, flags)) - }; - let (name, flags) = parser.parse2(tokens).unwrap(); - return Some(ItemFlag { name, flags }); - } - 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) { - Cow::Owned(format!("%{name}")) - } else { - Cow::Borrowed(name) - } -} -// from https://docs.rs/wit-component/latest/src/wit_component/printing.rs.html#155-192 -fn is_keyword(name: &str) -> bool { - matches!( - name, - "use" - | "type" - | "func" - | "u8" - | "u16" - | "u32" - | "u64" - | "s8" - | "s16" - | "s32" - | "s64" - | "f32" - | "f64" - | "float32" - | "float64" - | "char" - | "resource" - | "record" - | "flags" - | "variant" - | "enum" - | "bool" - | "string" - | "option" - | "result" - | "future" - | "stream" - | "list" - | "own" - | "borrow" - | "_" - | "as" - | "from" - | "static" - | "interface" - | "tuple" - | "world" - | "import" - | "export" - | "package" - | "with" - | "include" - | "constructor" - | "error-context" - | "async" - ) -} - -/// Extract WIT definitions from a wasm component and write them to `out_dir`. -/// Equivalent to `wasm-tools component wit --out-dir `. -pub fn extract_wit(wasm_file: &Path, out_dir: &Path) -> Result<()> { - let wasm = std::fs::read(wasm_file) - .with_context(|| format!("failed to read {}", wasm_file.display()))?; - let decoded = wit_component::decode(&wasm) - .with_context(|| format!("failed to decode component {}", wasm_file.display()))?; - let resolve = decoded.resolve(); - let main_pkg = decoded.package(); - - std::fs::create_dir_all(out_dir)?; - - let mut names: HashMap<&str, HashMap<&str, usize>> = HashMap::new(); - for (_id, pkg) in resolve.packages.iter() { - *names - .entry(&pkg.name.name) - .or_default() - .entry(&pkg.name.namespace) - .or_insert(0) += 1; - } - - for (id, pkg) in resolve.packages.iter() { - let is_main = id == main_pkg; - let mut printer = WitPrinter::default(); - printer.print(resolve, id, &[])?; - let output = printer.output.to_string(); - let dir = if is_main { - out_dir.to_path_buf() - } else { - out_dir.join("deps") - }; - let packages_with_same_name = &names[pkg.name.name.as_str()]; - let packages_with_same_namespace = packages_with_same_name[pkg.name.namespace.as_str()]; - let stem = if packages_with_same_name.len() == 1 { - if packages_with_same_namespace == 1 { - pkg.name.name.clone() - } else { - pkg.name - .version - .as_ref() - .map(|ver| format!("{}@{}", pkg.name.name, ver)) - .unwrap_or_else(|| pkg.name.name.clone()) - } - } else if packages_with_same_namespace == 1 { - format!("{}:{}", pkg.name.namespace, pkg.name.name) - } else { - pkg.name.to_string() - }; - std::fs::create_dir_all(&dir)?; - let path = dir.join(format!("{stem}.wit")); - std::fs::write(&path, &output)?; - } - Ok(()) -} - -/// Generate Rust bindings from WIT definitions. -/// Equivalent to `wit-bindgen rust --world --generate-all --merge-structurally-equal-types=true --out-dir `. -pub fn generate_bindings(wit_dir: &Path, world_name: &str, out_dir: &Path) -> Result<()> { - // wit-bindgen may depend on a different wit-parser version than wit-component - use wit_bindgen_core::{Files, WorldGenerator, wit_parser::Resolve}; - let mut resolve = Resolve::default(); - let (pkg, _) = resolve.push_dir(wit_dir)?; - let world = resolve.select_world(&[pkg], Some(world_name))?; - - let opts = wit_bindgen_rust::Opts { - generate_all: true, - merge_structurally_equal_types: Some(Some(true)), - ..Default::default() - }; - let mut generator = opts.build(); - let mut files = Files::default(); - generator.generate(&mut resolve, world, &mut files)?; - - std::fs::create_dir_all(out_dir)?; - for (name, content) in files.iter() { - let path = out_dir.join(name); - std::fs::write(&path, content)?; - } - Ok(()) -} - -/// Return the set of package ids that will have version suffix in the wit-bindgen -/// Use the same logic as in https://github.com/bytecodealliance/wit-bindgen/blob/main/crates/core/src/path.rs -pub fn package_with_version(resolve: &Resolve) -> BTreeSet { - let mut seen: BTreeMap<(String, String), Vec> = BTreeMap::new(); - for (id, p) in resolve.packages.iter() { - seen.entry((p.name.namespace.clone(), p.name.name.clone())) - .or_default() - .push(id); - } - let mut res = BTreeSet::new(); - for (_, ids) in seen { - if ids.len() > 1 { - res.extend(ids); - } - } - res -} diff --git a/src/wit_util.rs b/src/wit_util.rs new file mode 100644 index 0000000..6b97015 --- /dev/null +++ b/src/wit_util.rs @@ -0,0 +1,178 @@ +use anyhow::{Context, Result}; +use std::borrow::Cow; +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::path::Path; +use wit_component::WitPrinter; +use wit_parser::{PackageId, Resolve}; + +// 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 +// (see `crate::codegen::util::is_wrapped_module`). +const WRAPPED_WIT: &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) +} + +// utils for WIT names +pub fn ident(name: &str) -> Cow<'_, str> { + if is_keyword(name) { + Cow::Owned(format!("%{name}")) + } else { + Cow::Borrowed(name) + } +} +// from https://docs.rs/wit-component/latest/src/wit_component/printing.rs.html#155-192 +fn is_keyword(name: &str) -> bool { + matches!( + name, + "use" + | "type" + | "func" + | "u8" + | "u16" + | "u32" + | "u64" + | "s8" + | "s16" + | "s32" + | "s64" + | "f32" + | "f64" + | "float32" + | "float64" + | "char" + | "resource" + | "record" + | "flags" + | "variant" + | "enum" + | "bool" + | "string" + | "option" + | "result" + | "future" + | "stream" + | "list" + | "own" + | "borrow" + | "_" + | "as" + | "from" + | "static" + | "interface" + | "tuple" + | "world" + | "import" + | "export" + | "package" + | "with" + | "include" + | "constructor" + | "error-context" + | "async" + ) +} + +/// Extract WIT definitions from a wasm component and write them to `out_dir`. +/// Equivalent to `wasm-tools component wit --out-dir `. +pub fn extract_wit(wasm_file: &Path, out_dir: &Path) -> Result<()> { + let wasm = std::fs::read(wasm_file) + .with_context(|| format!("failed to read {}", wasm_file.display()))?; + let decoded = wit_component::decode(&wasm) + .with_context(|| format!("failed to decode component {}", wasm_file.display()))?; + let resolve = decoded.resolve(); + let main_pkg = decoded.package(); + + std::fs::create_dir_all(out_dir)?; + + let mut names: HashMap<&str, HashMap<&str, usize>> = HashMap::new(); + for (_id, pkg) in resolve.packages.iter() { + *names + .entry(&pkg.name.name) + .or_default() + .entry(&pkg.name.namespace) + .or_insert(0) += 1; + } + + for (id, pkg) in resolve.packages.iter() { + let is_main = id == main_pkg; + let mut printer = WitPrinter::default(); + printer.print(resolve, id, &[])?; + let output = printer.output.to_string(); + let dir = if is_main { + out_dir.to_path_buf() + } else { + out_dir.join("deps") + }; + let packages_with_same_name = &names[pkg.name.name.as_str()]; + let packages_with_same_namespace = packages_with_same_name[pkg.name.namespace.as_str()]; + let stem = if packages_with_same_name.len() == 1 { + if packages_with_same_namespace == 1 { + pkg.name.name.clone() + } else { + pkg.name + .version + .as_ref() + .map(|ver| format!("{}@{}", pkg.name.name, ver)) + .unwrap_or_else(|| pkg.name.name.clone()) + } + } else if packages_with_same_namespace == 1 { + format!("{}:{}", pkg.name.namespace, pkg.name.name) + } else { + pkg.name.to_string() + }; + std::fs::create_dir_all(&dir)?; + let path = dir.join(format!("{stem}.wit")); + std::fs::write(&path, &output)?; + } + Ok(()) +} + +/// Generate Rust bindings from WIT definitions. +/// Equivalent to `wit-bindgen rust --world --generate-all --merge-structurally-equal-types=true --out-dir `. +pub fn generate_bindings(wit_dir: &Path, world_name: &str, out_dir: &Path) -> Result<()> { + // wit-bindgen may depend on a different wit-parser version than wit-component + use wit_bindgen_core::{Files, WorldGenerator, wit_parser::Resolve}; + let mut resolve = Resolve::default(); + let (pkg, _) = resolve.push_dir(wit_dir)?; + let world = resolve.select_world(&[pkg], Some(world_name))?; + + let opts = wit_bindgen_rust::Opts { + generate_all: true, + merge_structurally_equal_types: Some(Some(true)), + ..Default::default() + }; + let mut generator = opts.build(); + let mut files = Files::default(); + generator.generate(&mut resolve, world, &mut files)?; + + std::fs::create_dir_all(out_dir)?; + for (name, content) in files.iter() { + let path = out_dir.join(name); + std::fs::write(&path, content)?; + } + Ok(()) +} + +/// Return the set of package ids that will have version suffix in the wit-bindgen +/// Use the same logic as in https://github.com/bytecodealliance/wit-bindgen/blob/main/crates/core/src/path.rs +pub fn package_with_version(resolve: &Resolve) -> BTreeSet { + let mut seen: BTreeMap<(String, String), Vec> = BTreeMap::new(); + for (id, p) in resolve.packages.iter() { + seen.entry((p.name.namespace.clone(), p.name.name.clone())) + .or_default() + .push(id); + } + let mut res = BTreeSet::new(); + for (_, ids) in seen { + if ids.len() > 1 { + res.extend(ids); + } + } + res +} From b8d30985b12d71614b42ff209a6b8bdc1ae321ff Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 20:17:17 -0700 Subject: [PATCH 2/4] split out start and export in codegen --- src/codegen/dialog.rs | 184 +++++++++++++++++------------------------- src/codegen/fuzz.rs | 173 ++++++++++++++++----------------------- src/codegen/mod.rs | 88 +++++++++++++++++--- src/codegen/record.rs | 9 +-- src/codegen/replay.rs | 171 ++++++++++++++++----------------------- src/codegen/util.rs | 27 +++++++ 6 files changed, 324 insertions(+), 328 deletions(-) diff --git a/src/codegen/dialog.rs b/src/codegen/dialog.rs index ee0dfe3..ad62c0b 100644 --- a/src/codegen/dialog.rs +++ b/src/codegen/dialog.rs @@ -1,129 +1,97 @@ -use super::State; use super::util::{ - FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, - get_return_type, make_path, wit_func_name, + constructor_resource_name, extract_arg_info, get_return_type, params_to_wave, wit_func_name, }; +use super::{ExportFunc, State}; use quote::quote; -use syn::{Signature, parse_quote, visit_mut::VisitMut}; +use syn::{Signature, parse_quote}; impl State { - pub fn generate_dialog_func( + pub fn generate_dialog_import_func( &self, module_path: &[String], sig: &Signature, resource: &Option, ) -> syn::ImplItemFn { let func_name = &sig.ident; - let is_export = module_path.join("::") == "exports::proxy::recorder::start_replay"; - if !is_export { - let (kind, args) = extract_arg_info(sig); - let arg_names = args.iter().map(|arg| &arg.ident); - let display_name = wit_func_name(module_path, resource, func_name, &kind); - let ret_ty = get_return_type(&sig.output); - if let Some(ty) = ret_ty { - let init_vec = if matches!(kind, Some(ResourceFuncKind::Method)) { - quote! { vec![wasm_wave::to_string(&ToValue::to_value(&self)).unwrap()] } - } else { - quote! { Vec::new() } - }; - let (ret_name, read_ret) = match constructor_resource_name(resource, &kind) { - Some(name) => ( - quote! { #name }, - quote! { MockedResource { name: #name.to_string(), ..Dialog::read_value(0) } }, - ), - None => ( - quote! { <#ty as WitName>::name() }, - quote! { Dialog::read_value(0) }, - ), - }; - parse_quote! { - #sig { - let mut __params: Vec = #init_vec; - #( - __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_names)).unwrap()); - )* - proxy::util::dialog::print(0, &format!("import: {}({})", #display_name, __params.join(", "))); - proxy::util::dialog::print(0, &format!("return type: {}", #ret_name)); - let ret = #read_ret; - proxy::util::dialog::print(0, &format!("ret: {}", wasm_wave::to_string(&ToValue::to_value(&ret)).unwrap())); - ret - } - } - } else { - parse_quote! { - #[allow(unused_variables)] - #sig {} + let (kind, args) = extract_arg_info(sig); + let display_name = wit_func_name(module_path, resource, func_name, &kind); + let ret_ty = get_return_type(&sig.output); + if let Some(ty) = ret_ty { + let params = params_to_wave(&kind, &args); + let (ret_name, read_ret) = match constructor_resource_name(resource, &kind) { + Some(name) => ( + quote! { #name }, + quote! { MockedResource { name: #name.to_string(), ..Dialog::read_value(0) } }, + ), + None => ( + quote! { <#ty as WitName>::name() }, + quote! { Dialog::read_value(0) }, + ), + }; + parse_quote! { + #sig { + #params + proxy::util::dialog::print(0, &format!("import: {}({})", #display_name, __params.join(", "))); + proxy::util::dialog::print(0, &format!("return type: {}", #ret_name)); + let ret = #read_ret; + proxy::util::dialog::print(0, &format!("ret: {}", wasm_wave::to_string(&ToValue::to_value(&ret)).unwrap())); + ret } } } else { - assert!(func_name == "start"); - let arms: Vec<_> = self - .funcs - .iter() - .filter(|(path, _)| path[0] != "exports" && path[0] != "proxy") - .flat_map(|(path, resources)| { - resources.iter().flat_map(move |(resource, sigs)| { - sigs.iter().filter_map(move |sig| { - let (kind, args) = extract_arg_info(sig); - if matches!(kind, Some(ResourceFuncKind::Method)) { - return None; - } - let arg_name: Vec<_> = args.iter().map(|arg| &arg.ident).collect(); - let call_param = args.iter().map(|arg| arg.call_param()); - let ty = args.iter().map(|arg| { - let mut ty = arg.ty.clone(); - FullTypePath { module_path: path }.visit_type_mut(&mut ty); - if let Some(owned) = get_owned_type(&ty) { - owned - } else { - ty - } - }); - let func_name = if let Some(resource) = resource { - format!("{}::{}", resource, sig.ident) - } else { - sig.ident.to_string() - }; - let func = make_path(path, &func_name); - let display_name = wit_func_name(path, resource, &sig.ident, &kind); - Some((quote! { - { - proxy::util::dialog::print(0, &format!("call export func {}", #display_name)); - let mut __params: Vec = Vec::new(); - #( - proxy::util::dialog::print(0, &format!("provide argument for {}: {}", stringify!(#arg_name), <#ty as WitName>::name())); - let #arg_name: #ty = Dialog::read_value(0); - __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_name)).unwrap()); - )* - proxy::util::dialog::print(0, &format!("export: {}({})", #display_name, __params.join(", "))); - let _ = #func(#(#call_param),*); - } - }, display_name)) - }) - }) - }) - .collect(); - let (arms, display_names): (Vec<_>, Vec<_>) = arms.into_iter().unzip(); - let display_names = - quote! { ["All done".to_string(), #(#display_names.to_string()),*] }; - let func_len = arms.iter().len(); - let idxs = 1..=func_len; parse_quote! { - #sig { - loop { - let idx = proxy::util::dialog::read_select(0, "Select an export function to call", &#display_names) as usize; - match idx { - 0 => break, - #(#idxs => #arms)* - _ => unreachable!(), - } - // clean up borrowed resources from input args - SCOPED_ALLOC.with(|alloc| { - alloc.borrow_mut().clear(); - }); + #[allow(unused_variables)] + #sig {} + } + } + } + pub fn generate_dialog_start_func(&self, sig: &Signature) -> syn::ImplItemFn { + let funcs = self.export_funcs(); + let display_names = funcs.iter().map(|func| &func.display_name); + let display_names = quote! { ["All done".to_string(), #(#display_names.to_string()),*] }; + let arms: Vec<_> = funcs + .iter() + .map(|func| { + let ExportFunc { + display_name, + func, + arg_names, + arg_tys, + call_params, + .. + } = func; + quote! { + { + proxy::util::dialog::print(0, &format!("call export func {}", #display_name)); + let mut __params: Vec = Vec::new(); + #( + proxy::util::dialog::print(0, &format!("provide argument for {}: {}", stringify!(#arg_names), <#arg_tys as WitName>::name())); + let #arg_names: #arg_tys = Dialog::read_value(0); + __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_names)).unwrap()); + )* + proxy::util::dialog::print(0, &format!("export: {}({})", #display_name, __params.join(", "))); + let _ = #func(#(#call_params),*); + } } + }) + .collect(); + let func_len = arms.len(); + let idxs = 1..=func_len; + parse_quote! { + #sig { + loop { + let idx = proxy::util::dialog::read_select(0, "Select an export function to call", &#display_names) as usize; + match idx { + 0 => break, + #(#idxs => #arms)* + _ => unreachable!(), } + // clean up borrowed resources from input args + SCOPED_ALLOC.with(|alloc| { + alloc.borrow_mut().clear(); + }); } + } } } } diff --git a/src/codegen/fuzz.rs b/src/codegen/fuzz.rs index 33849b9..528178c 100644 --- a/src/codegen/fuzz.rs +++ b/src/codegen/fuzz.rs @@ -1,122 +1,91 @@ -use super::State; use super::util::{ - FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, - get_return_type, make_path, wit_func_name, + constructor_resource_name, extract_arg_info, get_return_type, params_to_wave, wit_func_name, }; +use super::{ExportFunc, State}; use quote::quote; -use syn::{Signature, parse_quote, visit_mut::VisitMut}; +use syn::{Signature, parse_quote}; impl State { - pub fn generate_fuzz_func( + pub fn generate_fuzz_import_func( &self, module_path: &[String], sig: &Signature, resource: &Option, ) -> syn::ImplItemFn { let func_name = &sig.ident; - let is_export = module_path.join("::") == "exports::proxy::recorder::start_replay"; - if !is_export { - let (kind, args) = extract_arg_info(sig); - let arg_names = args.iter().map(|arg| &arg.ident); - let display_name = wit_func_name(module_path, resource, func_name, &kind); - let ret_ty = get_return_type(&sig.output); - if ret_ty.is_some() { - let init_vec = if matches!(kind, Some(ResourceFuncKind::Method)) { - quote! { vec![wasm_wave::to_string(&ToValue::to_value(&self)).unwrap()] } - } else { - quote! { Vec::new() } - }; - let gen_ret = match constructor_resource_name(resource, &kind) { - Some(name) => { - quote! { MockedResource { name: #name.to_string(), ..u.arbitrary().unwrap() } } - } - None => quote! { u.arbitrary().unwrap() }, - }; - parse_quote! { - #sig { - let mut __params: Vec = #init_vec; - #( - __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_names)).unwrap()); - )* - let mut __buf = __params.join(","); - proxy::util::debug::print(&format!("import: {}({})", #display_name, __buf)); - __buf += #display_name; - let mut u = Unstructured::new(&__buf.as_bytes()); - let res = #gen_ret; - let res_str = wasm_wave::to_string(&ToValue::to_value(&res)).unwrap(); - proxy::util::debug::print(&format!("ret: {}", res_str)); - res - } + let (kind, args) = extract_arg_info(sig); + let display_name = wit_func_name(module_path, resource, func_name, &kind); + let ret_ty = get_return_type(&sig.output); + if ret_ty.is_some() { + let params = params_to_wave(&kind, &args); + let gen_ret = match constructor_resource_name(resource, &kind) { + Some(name) => { + quote! { MockedResource { name: #name.to_string(), ..u.arbitrary().unwrap() } } } - } else { - parse_quote! { - #[allow(unused_variables)] - #sig {} + None => quote! { u.arbitrary().unwrap() }, + }; + parse_quote! { + #sig { + #params + let mut __buf = __params.join(","); + proxy::util::debug::print(&format!("import: {}({})", #display_name, __buf)); + __buf += #display_name; + let mut u = Unstructured::new(&__buf.as_bytes()); + let res = #gen_ret; + let res_str = wasm_wave::to_string(&ToValue::to_value(&res)).unwrap(); + proxy::util::debug::print(&format!("ret: {}", res_str)); + res } } } else { - assert!(func_name == "start"); - let arms: Vec<_> = self - .funcs - .iter() - .filter(|(path, _)| path[0] != "exports" && path[0] != "proxy") - .flat_map(|(path, resources)| { - resources.iter().flat_map(move |(resource, sigs)| { - sigs.iter().filter_map(move |sig| { - let (kind, args) = extract_arg_info(sig); - if matches!(kind, Some(ResourceFuncKind::Method)) { - return None; - } - let arg_name: Vec<_> = args.iter().map(|arg| &arg.ident).collect(); - let call_param = args.iter().map(|arg| arg.call_param()); - let ty = args.iter().map(|arg| { - let mut ty = arg.ty.clone(); - FullTypePath { module_path: path }.visit_type_mut(&mut ty); - if let Some(owned) = get_owned_type(&ty) { - owned - } else { - ty - } - }); - let func_name = if let Some(resource) = resource { - format!("{}::{}", resource, sig.ident) - } else { - sig.ident.to_string() - }; - let func = make_path(path, &func_name); - let display_name = wit_func_name(path, resource, &sig.ident, &kind); - Some(quote! { - { - let mut __params: Vec = Vec::new(); - #( - let #arg_name: #ty = u.arbitrary().unwrap(); - __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_name)).unwrap()); - )* - proxy::util::debug::print(&format!("export: {}({})", #display_name, __params.join(", "))); - let _ = #func(#(#call_param),*); - } - }) - }) - }) - }) - .collect(); - let func_len = arms.iter().len(); - let idxs = 1..=func_len; parse_quote! { - #sig { - let __buf = proxy::util::debug::get_random(); - let mut u = Unstructured::new(&__buf); - for _ in 0..10 { - let idx = u.int_in_range(1..=#func_len).unwrap(); - match idx { - #(#idxs => #arms)* - _ => unreachable!(), - } - // clean up borrowed resources from input args - SCOPED_ALLOC.with(|alloc| { - alloc.borrow_mut().clear(); - }); + #[allow(unused_variables)] + #sig {} + } + } + } + pub fn generate_fuzz_start_func(&self, sig: &Signature) -> syn::ImplItemFn { + let arms: Vec<_> = self + .export_funcs() + .into_iter() + .map(|func| { + let ExportFunc { + display_name, + func, + arg_names, + arg_tys, + call_params, + .. + } = func; + quote! { + { + let mut __params: Vec = Vec::new(); + #( + let #arg_names: #arg_tys = u.arbitrary().unwrap(); + __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_names)).unwrap()); + )* + proxy::util::debug::print(&format!("export: {}({})", #display_name, __params.join(", "))); + let _ = #func(#(#call_params),*); + } + } + }) + .collect(); + let func_len = arms.len(); + let idxs = 1..=func_len; + parse_quote! { + #sig { + let __buf = proxy::util::debug::get_random(); + let mut u = Unstructured::new(&__buf); + for _ in 0..10 { + let idx = u.int_in_range(1..=#func_len).unwrap(); + match idx { + #(#idxs => #arms)* + _ => unreachable!(), } + // clean up borrowed resources from input args + SCOPED_ALLOC.with(|alloc| { + alloc.borrow_mut().clear(); + }); } } } diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index caecd8b..9573af4 100644 --- a/src/codegen/mod.rs +++ b/src/codegen/mod.rs @@ -8,7 +8,8 @@ use syn::{ visit_mut::VisitMut, }; use util::{ - FullTypePath, get_resource_from_trait_name, get_return_type, make_path, toggle_wrapped_module, + FullTypePath, ResourceFuncKind, extract_arg_info, get_owned_type, get_resource_from_trait_name, + get_return_type, make_path, toggle_wrapped_module, wit_func_name, }; mod dialog; @@ -72,6 +73,24 @@ pub struct ItemFlag { pub name: Ident, pub flags: Vec, } +/// An exported function of the main component, see `State::export_funcs`. +struct ExportFunc { + /// WIT name, e.g. `ns:pkg/iface.func` + display_name: String, + /// Rust path to the function + func: syn::Path, + arg_names: Vec, + /// Owned argument types, e.g. `String` for `&str` + arg_tys: Vec, + /// Arguments at the call site, borrowed where the function expects a reference + call_params: Vec, + has_ret: bool, +} + +/// Interface for converting resources between the generated components. +const CONVERSION_MODULE: &str = "exports::proxy::conversion::conversion"; +/// Interface whose `start` function drives the main component's exports in virtualized modes. +const START_REPLAY_MODULE: &str = "exports::proxy::recorder::start_replay"; impl GenerateArgs { pub fn generate(&self) -> Result<()> { @@ -167,28 +186,37 @@ impl State { let mut transformer = FullTypePath { module_path }; transformer.visit_signature_mut(&mut sig); let module_full_name = module_path.join("::"); - if module_full_name == "exports::proxy::conversion::conversion" { + if module_full_name == CONVERSION_MODULE { let stub_impl = self.generate_conversion_func(&sig); methods.push(syn::ImplItem::Fn(stub_impl)); continue; } - let stub_impl = match self.mode { - GenerateMode::Stubs => parse_quote! { + let is_start = module_full_name == START_REPLAY_MODULE; + if is_start { + assert!(sig.ident == "start"); + } + let stub_impl = match (&self.mode, is_start) { + (GenerateMode::Stubs, _) => parse_quote! { #[allow(unused_variables)] #sig { unimplemented!() } }, - GenerateMode::Instrument | GenerateMode::Record => { + (GenerateMode::Instrument | GenerateMode::Record, _) => { self.generate_instrument_func(module_path, &sig, &resource) } - GenerateMode::Replay => { - self.generate_replay_func(module_path, &sig, &resource) + (GenerateMode::Replay, false) => { + self.generate_replay_import_func(module_path, &sig, &resource) } - GenerateMode::Fuzz => self.generate_fuzz_func(module_path, &sig, &resource), - GenerateMode::Dialog => { - self.generate_dialog_func(module_path, &sig, &resource) + (GenerateMode::Replay, true) => self.generate_replay_start_func(&sig), + (GenerateMode::Fuzz, false) => { + self.generate_fuzz_import_func(module_path, &sig, &resource) } + (GenerateMode::Fuzz, true) => self.generate_fuzz_start_func(&sig), + (GenerateMode::Dialog, false) => { + self.generate_dialog_import_func(module_path, &sig, &resource) + } + (GenerateMode::Dialog, true) => self.generate_dialog_start_func(&sig), }; methods.push(syn::ImplItem::Fn(stub_impl)); } @@ -246,6 +274,46 @@ impl State { .to_string() .to_kebab_case() } + /// Functions exported by the main component, which the `start` function of the virtualized + /// modes calls. Resource methods are skipped. + fn export_funcs(&self) -> Vec { + let mut res = Vec::new(); + for (path, resources) in self + .funcs + .iter() + .filter(|(path, _)| path[0] != "exports" && path[0] != "proxy") + { + for (resource, sigs) in resources { + for sig in sigs { + let (kind, args) = extract_arg_info(sig); + if matches!(kind, Some(ResourceFuncKind::Method)) { + continue; + } + let arg_tys = args + .iter() + .map(|arg| { + let mut ty = arg.ty.clone(); + FullTypePath { module_path: path }.visit_type_mut(&mut ty); + get_owned_type(&ty).unwrap_or(ty) + }) + .collect(); + let func_name = match resource { + Some(resource) => format!("{}::{}", resource, sig.ident), + None => sig.ident.to_string(), + }; + res.push(ExportFunc { + display_name: wit_func_name(path, resource, &sig.ident, &kind), + func: make_path(path, &func_name), + arg_names: args.iter().map(|arg| arg.ident.clone()).collect(), + arg_tys, + call_params: args.iter().map(|arg| arg.call_param()).collect(), + has_ret: get_return_type(&sig.output).is_some(), + }); + } + } + } + res + } fn into_output_file(self) -> File { File { items: self.output, diff --git a/src/codegen/record.rs b/src/codegen/record.rs index 5bdbfaf..a303942 100644 --- a/src/codegen/record.rs +++ b/src/codegen/record.rs @@ -1,5 +1,6 @@ use super::util::{ - ResourceFuncKind, extract_arg_info, get_return_type, is_wrapped_module, wit_func_name, + ResourceFuncKind, extract_arg_info, get_return_type, init_params, is_wrapped_module, + wit_func_name, }; use super::{GenerateMode, State, get_proxy_path}; use quote::quote; @@ -52,11 +53,7 @@ impl State { } }, GenerateMode::Record => { - let init_vec = if matches!(kind, Some(ResourceFuncKind::Method)) { - quote! { vec![wasm_wave::to_string(&ToValue::to_value(&self)).unwrap()] } - } else { - quote! { Vec::new() } - }; + let init_vec = init_params(&kind); let is_mut = if args.is_empty() { quote! {} } else { diff --git a/src/codegen/replay.rs b/src/codegen/replay.rs index ae59f24..a2027cf 100644 --- a/src/codegen/replay.rs +++ b/src/codegen/replay.rs @@ -1,125 +1,92 @@ -use super::State; use super::util::{ - FullTypePath, ResourceFuncKind, constructor_resource_name, extract_arg_info, get_owned_type, - get_return_type, make_path, wit_func_name, + constructor_resource_name, extract_arg_info, get_return_type, self_wave, wit_func_name, }; +use super::{ExportFunc, State}; use quote::quote; -use syn::{Signature, parse_quote, visit_mut::VisitMut}; +use syn::{Signature, parse_quote}; impl State { - pub fn generate_replay_func( + pub fn generate_replay_import_func( &self, module_path: &[String], sig: &Signature, resource: &Option, ) -> syn::ImplItemFn { let func_name = &sig.ident; - let is_export = module_path.join("::") == "exports::proxy::recorder::start_replay"; - if !is_export { - let (kind, args) = extract_arg_info(sig); - let arg_names = args.iter().map(|arg| &arg.ident); - let display_name = wit_func_name(module_path, resource, func_name, &kind); - let ret_ty = get_return_type(&sig.output); - let replay_import = if let Some(ret_ty) = ret_ty { - let to_rust = match constructor_resource_name(resource, &kind) { - Some(name) => { - quote! { MockedResource { name: #name.to_string(), ..ret.to_rust() } } - } - None => quote! { ret.to_rust() }, - }; + let (kind, args) = extract_arg_info(sig); + let arg_names = args.iter().map(|arg| &arg.ident); + let display_name = wit_func_name(module_path, resource, func_name, &kind); + let ret_ty = get_return_type(&sig.output); + let replay_import = if let Some(ret_ty) = ret_ty { + let to_rust = match constructor_resource_name(resource, &kind) { + Some(name) => { + quote! { MockedResource { name: #name.to_string(), ..ret.to_rust() } } + } + None => quote! { ret.to_rust() }, + }; + quote! { + let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)).unwrap(); + let ret: Value = wasm_wave::from_str(&<#ret_ty as ValueTyped>::value_type(), &wave).unwrap(); + #to_rust + } + } else { + quote! { + let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)); + assert!(wave.is_none()); + } + }; + let self_value = self_wave(&kind).map(|self_wave| quote! { #self_wave, }); + parse_quote! { + #sig { + let args = vec![#self_value #( wasm_wave::to_string(&#arg_names.to_value()).unwrap() ),*]; + #replay_import + } + } + } + pub fn generate_replay_start_func(&self, sig: &Signature) -> syn::ImplItemFn { + let arms = self.export_funcs().into_iter().map(|func| { + let ExportFunc { + display_name, + func, + arg_names, + arg_tys, + call_params, + has_ret, + } = func; + let arg_idx = (0..arg_names.len()).map(|idx| quote! { args[#idx] }); + let assert_ret = if has_ret { quote! { - let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)).unwrap(); - let ret: Value = wasm_wave::from_str(&<#ret_ty as ValueTyped>::value_type(), &wave).unwrap(); - #to_rust + let wave_res = wasm_wave::to_string(&res.to_value()).unwrap(); + proxy::recorder::replay::assert_export_ret(Some(#display_name), Some(&wave_res)); } } else { quote! { - let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)); - assert!(wave.is_none()); + assert!(res == ()); + proxy::recorder::replay::assert_export_ret(Some(#display_name), None); } }; - let self_value = if matches!(kind, Some(ResourceFuncKind::Method)) { - // Use ToValue::to_value to avoid the auto-deref from self.to_value() - quote! { wasm_wave::to_string(&ToValue::to_value(&self)).unwrap(), } - } else { - quote! {} - }; - parse_quote! { - #sig { - let args = vec![#self_value #( wasm_wave::to_string(&#arg_names.to_value()).unwrap() ),*]; - #replay_import + quote! { + #display_name => { + #( + let arg_value: Value = wasm_wave::from_str(&<#arg_tys as ValueTyped>::value_type(), &#arg_idx).unwrap(); + let #arg_names: #arg_tys = arg_value.to_rust(); + )* + let res = #func(#(#call_params),*); + #assert_ret } } - } else { - assert!(func_name == "start"); - let arms = self - .funcs - .iter() - .filter(|(path, _)| path[0] != "exports" && path[0] != "proxy") - .flat_map(|(path, resources)| { - resources.iter().flat_map(move |(resource, sigs)| { - sigs.iter().filter_map(move |sig| { - let (kind, args) = extract_arg_info(sig); - if matches!(kind, Some(ResourceFuncKind::Method)) { - return None; - } - let arg_name: Vec<_> = args.iter().map(|arg| &arg.ident).collect(); - let arg_idx = args.iter().enumerate().map(|(idx, _)| quote! { args[#idx] }); - let call_param = args.iter().map(|arg| arg.call_param()); - let ty = args.iter().map(|arg| { - let mut ty = arg.ty.clone(); - FullTypePath { - module_path: path, - }.visit_type_mut(&mut ty); - if let Some(owned) = get_owned_type(&ty) { - owned - } else { - ty - } - }); - let func_name = if let Some(resource) = resource { - format!("{}::{}", resource, sig.ident) - } else { - sig.ident.to_string() - }; - let func = make_path(path, &func_name); - let display_name = wit_func_name(path, resource, &sig.ident, &kind); - let assert_ret = if get_return_type(&sig.output).is_none() { - quote! { - assert!(res == ()); - proxy::recorder::replay::assert_export_ret(Some(#display_name), None); - } - } else { - quote! { - let wave_res = wasm_wave::to_string(&res.to_value()).unwrap(); - proxy::recorder::replay::assert_export_ret(Some(#display_name), Some(&wave_res)); - } - }; - Some(quote! { - #display_name => { - #( - let arg_value: Value = wasm_wave::from_str(&<#ty as ValueTyped>::value_type(), &#arg_idx).unwrap(); - let #arg_name: #ty = arg_value.to_rust(); - )* - let res = #func(#(#call_param),*); - #assert_ret - } - }) - }) - }) - }); - parse_quote! { - #sig { - while let Some((method, args)) = proxy::recorder::replay::replay_export() { - match method.as_str() { - #(#arms)* - _ => unreachable!(), - } - // clean up borrowed resources from input args - SCOPED_ALLOC.with(|alloc| { - alloc.borrow_mut().clear(); - }); + }); + parse_quote! { + #sig { + while let Some((method, args)) = proxy::recorder::replay::replay_export() { + match method.as_str() { + #(#arms)* + _ => unreachable!(), } + // clean up borrowed resources from input args + SCOPED_ALLOC.with(|alloc| { + alloc.borrow_mut().clear(); + }); } } } diff --git a/src/codegen/util.rs b/src/codegen/util.rs index 280be9f..4069457 100644 --- a/src/codegen/util.rs +++ b/src/codegen/util.rs @@ -1,4 +1,6 @@ use heck::ToKebabCase; +use proc_macro2::TokenStream; +use quote::quote; use syn::{FnArg, Ident, Signature, Type, parse_quote, visit_mut::VisitMut}; pub struct FullTypePath<'a> { pub module_path: &'a [String], @@ -140,6 +142,31 @@ impl<'a> VisitMut for FullTypePath<'a> { } } +/// WAVE string of `self`, the first param of a resource method. +// Use ToValue::to_value to avoid the auto-deref from self.to_value() +pub fn self_wave(kind: &Option) -> Option { + matches!(kind, Some(ResourceFuncKind::Method)) + .then(|| quote! { wasm_wave::to_string(&ToValue::to_value(&self)).unwrap() }) +} +/// Initial `Vec` of WAVE params: `self` for methods, empty otherwise. +pub fn init_params(kind: &Option) -> TokenStream { + match self_wave(kind) { + Some(self_wave) => quote! { vec![#self_wave] }, + None => quote! { Vec::new() }, + } +} +/// Binds `__params` to the WAVE strings of `self` (for methods) and each arg. +pub fn params_to_wave(kind: &Option, args: &[ArgInfo]) -> TokenStream { + let init = init_params(kind); + let arg_names = args.iter().map(|arg| &arg.ident); + quote! { + let mut __params: Vec = #init; + #( + __params.push(wasm_wave::to_string(&ToValue::to_value(&#arg_names)).unwrap()); + )* + } +} + pub fn get_resource_from_trait_name(trait_name: &str) -> Option { let resource = trait_name.strip_prefix("Guest").unwrap(); match resource { From f1f0ea9d95bc7920edbdb07e79bd8319b165b05a Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 21:47:03 -0700 Subject: [PATCH 3/4] refactor replay param --- src/codegen/replay.rs | 11 +++++------ src/codegen/util.rs | 14 +++++--------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/src/codegen/replay.rs b/src/codegen/replay.rs index a2027cf..21e5ce4 100644 --- a/src/codegen/replay.rs +++ b/src/codegen/replay.rs @@ -1,5 +1,5 @@ use super::util::{ - constructor_resource_name, extract_arg_info, get_return_type, self_wave, wit_func_name, + constructor_resource_name, extract_arg_info, get_return_type, params_to_wave, wit_func_name, }; use super::{ExportFunc, State}; use quote::quote; @@ -14,7 +14,6 @@ impl State { ) -> syn::ImplItemFn { let func_name = &sig.ident; let (kind, args) = extract_arg_info(sig); - let arg_names = args.iter().map(|arg| &arg.ident); let display_name = wit_func_name(module_path, resource, func_name, &kind); let ret_ty = get_return_type(&sig.output); let replay_import = if let Some(ret_ty) = ret_ty { @@ -25,20 +24,20 @@ impl State { None => quote! { ret.to_rust() }, }; quote! { - let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)).unwrap(); + let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&__params)).unwrap(); let ret: Value = wasm_wave::from_str(&<#ret_ty as ValueTyped>::value_type(), &wave).unwrap(); #to_rust } } else { quote! { - let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&args)); + let wave = proxy::recorder::replay::replay_import(Some(#display_name), Some(&__params)); assert!(wave.is_none()); } }; - let self_value = self_wave(&kind).map(|self_wave| quote! { #self_wave, }); + let params = params_to_wave(&kind, &args); parse_quote! { #sig { - let args = vec![#self_value #( wasm_wave::to_string(&#arg_names.to_value()).unwrap() ),*]; + #params #replay_import } } diff --git a/src/codegen/util.rs b/src/codegen/util.rs index 4069457..4220d54 100644 --- a/src/codegen/util.rs +++ b/src/codegen/util.rs @@ -142,17 +142,13 @@ impl<'a> VisitMut for FullTypePath<'a> { } } -/// WAVE string of `self`, the first param of a resource method. -// Use ToValue::to_value to avoid the auto-deref from self.to_value() -pub fn self_wave(kind: &Option) -> Option { - matches!(kind, Some(ResourceFuncKind::Method)) - .then(|| quote! { wasm_wave::to_string(&ToValue::to_value(&self)).unwrap() }) -} /// Initial `Vec` of WAVE params: `self` for methods, empty otherwise. pub fn init_params(kind: &Option) -> TokenStream { - match self_wave(kind) { - Some(self_wave) => quote! { vec![#self_wave] }, - None => quote! { Vec::new() }, + if matches!(kind, Some(ResourceFuncKind::Method)) { + // Use ToValue::to_value to avoid the auto-deref from self.to_value() + quote! { vec![wasm_wave::to_string(&ToValue::to_value(&self)).unwrap()] } + } else { + quote! { Vec::new() } } } /// Binds `__params` to the WAVE strings of `self` (for methods) and each arg. From e15fd911b2c2d40366362f035f6569928a46ccb9 Mon Sep 17 00:00:00 2001 From: Yan Chen Date: Sun, 4 Oct 2026 22:10:07 -0700 Subject: [PATCH 4/4] fix --- src/codegen/dialog.rs | 8 +-- src/codegen/fuzz.rs | 8 +-- src/codegen/mod.rs | 40 ++++++++------ src/codegen/parse.rs | 4 +- src/codegen/record.rs | 6 +- src/codegen/replay.rs | 8 +-- src/codegen/util.rs | 14 ++--- src/run.rs | 125 +++++++++++++++++------------------------- src/traits/dialog.rs | 6 +- src/traits/wit.rs | 6 +- 10 files changed, 104 insertions(+), 121 deletions(-) diff --git a/src/codegen/dialog.rs b/src/codegen/dialog.rs index ad62c0b..7269d55 100644 --- a/src/codegen/dialog.rs +++ b/src/codegen/dialog.rs @@ -10,15 +10,15 @@ impl State { &self, module_path: &[String], sig: &Signature, - resource: &Option, + resource: Option<&str>, ) -> syn::ImplItemFn { let func_name = &sig.ident; let (kind, args) = extract_arg_info(sig); - let display_name = wit_func_name(module_path, resource, func_name, &kind); + let display_name = wit_func_name(module_path, resource, func_name, kind.as_ref()); let ret_ty = get_return_type(&sig.output); if let Some(ty) = ret_ty { - let params = params_to_wave(&kind, &args); - let (ret_name, read_ret) = match constructor_resource_name(resource, &kind) { + let params = params_to_wave(kind.as_ref(), &args); + let (ret_name, read_ret) = match constructor_resource_name(resource, kind.as_ref()) { Some(name) => ( quote! { #name }, quote! { MockedResource { name: #name.to_string(), ..Dialog::read_value(0) } }, diff --git a/src/codegen/fuzz.rs b/src/codegen/fuzz.rs index 528178c..90bcaf0 100644 --- a/src/codegen/fuzz.rs +++ b/src/codegen/fuzz.rs @@ -10,15 +10,15 @@ impl State { &self, module_path: &[String], sig: &Signature, - resource: &Option, + resource: Option<&str>, ) -> syn::ImplItemFn { let func_name = &sig.ident; let (kind, args) = extract_arg_info(sig); - let display_name = wit_func_name(module_path, resource, func_name, &kind); + let display_name = wit_func_name(module_path, resource, func_name, kind.as_ref()); let ret_ty = get_return_type(&sig.output); if ret_ty.is_some() { - let params = params_to_wave(&kind, &args); - let gen_ret = match constructor_resource_name(resource, &kind) { + let params = params_to_wave(kind.as_ref(), &args); + let gen_ret = match constructor_resource_name(resource, kind.as_ref()) { Some(name) => { quote! { MockedResource { name: #name.to_string(), ..u.arbitrary().unwrap() } } } diff --git a/src/codegen/mod.rs b/src/codegen/mod.rs index 9573af4..210a038 100644 --- a/src/codegen/mod.rs +++ b/src/codegen/mod.rs @@ -97,21 +97,11 @@ impl GenerateArgs { let file = std::fs::read_to_string(&self.bindings)?; let ast = syn::parse_file(&file)?; - let mut state = State { - mode: self.mode.clone(), - traits: BTreeMap::new(), - types: BTreeMap::new(), - funcs: BTreeMap::new(), - module_paths: BTreeSet::new(), - type_aliases: BTreeMap::new(), - output: Vec::new(), - }; + let mut state = State::new(self.mode.clone()); state.generate_preamble(); state.find_all_items(&ast.items, vec![]); state.generate_stubs(); - let trait_generator = crate::traits::TraitGenerator::new(&state); - let traits = trait_generator.generate(); - drop(trait_generator); + let traits = crate::traits::TraitGenerator::new(&state).generate(); state.output.extend(traits); let mut file = state.into_output_file(); if self.use_custom_allocator { @@ -129,6 +119,17 @@ impl GenerateArgs { } } impl State { + fn new(mode: GenerateMode) -> Self { + State { + mode, + traits: BTreeMap::new(), + types: BTreeMap::new(), + funcs: BTreeMap::new(), + module_paths: BTreeSet::new(), + type_aliases: BTreeMap::new(), + output: Vec::new(), + } + } fn generate_stubs(&mut self) { for (module_path, traits) in &self.traits { for trait_item in traits { @@ -203,18 +204,18 @@ impl State { } }, (GenerateMode::Instrument | GenerateMode::Record, _) => { - self.generate_instrument_func(module_path, &sig, &resource) + self.generate_instrument_func(module_path, &sig, resource.as_deref()) } (GenerateMode::Replay, false) => { - self.generate_replay_import_func(module_path, &sig, &resource) + self.generate_replay_import_func(module_path, &sig, resource.as_deref()) } (GenerateMode::Replay, true) => self.generate_replay_start_func(&sig), (GenerateMode::Fuzz, false) => { - self.generate_fuzz_import_func(module_path, &sig, &resource) + self.generate_fuzz_import_func(module_path, &sig, resource.as_deref()) } (GenerateMode::Fuzz, true) => self.generate_fuzz_start_func(&sig), (GenerateMode::Dialog, false) => { - self.generate_dialog_import_func(module_path, &sig, &resource) + self.generate_dialog_import_func(module_path, &sig, resource.as_deref()) } (GenerateMode::Dialog, true) => self.generate_dialog_start_func(&sig), }; @@ -302,7 +303,12 @@ impl State { None => sig.ident.to_string(), }; res.push(ExportFunc { - display_name: wit_func_name(path, resource, &sig.ident, &kind), + display_name: wit_func_name( + path, + resource.as_deref(), + &sig.ident, + kind.as_ref(), + ), func: make_path(path, &func_name), arg_names: args.iter().map(|arg| arg.ident.clone()).collect(), arg_tys, diff --git a/src/codegen/parse.rs b/src/codegen/parse.rs index 03668f7..d38e44b 100644 --- a/src/codegen/parse.rs +++ b/src/codegen/parse.rs @@ -131,11 +131,11 @@ impl State { pub fn find_function( &self, module_path: &[String], - resource: &Option, + resource: Option<&str>, func: &Ident, ) -> Option<&Signature> { let module = self.funcs.get(module_path)?; - let funcs = module.get(resource)?; + let funcs = module.get(&resource.map(str::to_string))?; funcs.iter().find(|sig| sig.ident == *func) } pub fn has_type_def(&self, module_path: &[String], name: &str) -> bool { diff --git a/src/codegen/record.rs b/src/codegen/record.rs index a303942..7c84b70 100644 --- a/src/codegen/record.rs +++ b/src/codegen/record.rs @@ -11,7 +11,7 @@ impl State { &self, module_path: &[String], sig: &Signature, - resource: &Option, + resource: Option<&str>, ) -> syn::ImplItemFn { let func_name = &sig.ident; let (kind, args) = extract_arg_info(sig); @@ -53,13 +53,13 @@ impl State { } }, GenerateMode::Record => { - let init_vec = init_params(&kind); + let init_vec = init_params(kind.as_ref()); let is_mut = if args.is_empty() { quote! {} } else { quote! { mut } }; - let display_name = wit_func_name(module_path, resource, func_name, &kind); + let display_name = wit_func_name(module_path, resource, func_name, kind.as_ref()); let is_export = !is_wrapped_module(&module_path[1]); let record_ret = if get_return_type(&sig.output).is_none() { quote! { diff --git a/src/codegen/replay.rs b/src/codegen/replay.rs index 21e5ce4..acbc049 100644 --- a/src/codegen/replay.rs +++ b/src/codegen/replay.rs @@ -10,14 +10,14 @@ impl State { &self, module_path: &[String], sig: &Signature, - resource: &Option, + resource: Option<&str>, ) -> syn::ImplItemFn { let func_name = &sig.ident; let (kind, args) = extract_arg_info(sig); - let display_name = wit_func_name(module_path, resource, func_name, &kind); + let display_name = wit_func_name(module_path, resource, func_name, kind.as_ref()); let ret_ty = get_return_type(&sig.output); let replay_import = if let Some(ret_ty) = ret_ty { - let to_rust = match constructor_resource_name(resource, &kind) { + let to_rust = match constructor_resource_name(resource, kind.as_ref()) { Some(name) => { quote! { MockedResource { name: #name.to_string(), ..ret.to_rust() } } } @@ -34,7 +34,7 @@ impl State { assert!(wave.is_none()); } }; - let params = params_to_wave(&kind, &args); + let params = params_to_wave(kind.as_ref(), &args); parse_quote! { #sig { #params diff --git a/src/codegen/util.rs b/src/codegen/util.rs index 4220d54..9dde2ac 100644 --- a/src/codegen/util.rs +++ b/src/codegen/util.rs @@ -30,9 +30,9 @@ pub fn make_path(module_path: &[String], name: &str) -> syn::Path { } pub fn wit_func_name( module_path: &[String], - resource: &Option, + resource: Option<&str>, func_name: &Ident, - kind: &Option, + kind: Option<&ResourceFuncKind>, ) -> String { assert!(module_path.len() >= 3); let mut module_path = module_path.to_vec(); @@ -62,11 +62,11 @@ pub fn wit_func_name( /// traits like `Dialog` or `Arbitrary` cannot recover the resource name. Returns the WIT name of /// the resource to patch into the `MockedResource`, if `kind` is a constructor. pub fn constructor_resource_name( - resource: &Option, - kind: &Option, + resource: Option<&str>, + kind: Option<&ResourceFuncKind>, ) -> Option { match kind { - Some(ResourceFuncKind::Constructor) => Some(resource.as_ref()?.to_kebab_case()), + Some(ResourceFuncKind::Constructor) => Some(resource?.to_kebab_case()), _ => None, } } @@ -143,7 +143,7 @@ impl<'a> VisitMut for FullTypePath<'a> { } /// Initial `Vec` of WAVE params: `self` for methods, empty otherwise. -pub fn init_params(kind: &Option) -> TokenStream { +pub fn init_params(kind: Option<&ResourceFuncKind>) -> TokenStream { if matches!(kind, Some(ResourceFuncKind::Method)) { // Use ToValue::to_value to avoid the auto-deref from self.to_value() quote! { vec![wasm_wave::to_string(&ToValue::to_value(&self)).unwrap()] } @@ -152,7 +152,7 @@ pub fn init_params(kind: &Option) -> TokenStream { } } /// Binds `__params` to the WAVE strings of `self` (for methods) and each arg. -pub fn params_to_wave(kind: &Option, args: &[ArgInfo]) -> TokenStream { +pub fn params_to_wave(kind: Option<&ResourceFuncKind>, args: &[ArgInfo]) -> TokenStream { let init = init_params(kind); let arg_names = args.iter().map(|arg| &arg.ident); quote! { diff --git a/src/run.rs b/src/run.rs index 3d967b2..46e8b81 100644 --- a/src/run.rs +++ b/src/run.rs @@ -1,7 +1,8 @@ +use anyhow::Context; use clap::Parser; use std::path::PathBuf; use trace::Logger; -use wasmtime::component::types::{ComponentFunc, ComponentItem as CItem}; +use wasmtime::component::types::{ComponentExtern, ComponentFunc, ComponentItem as CItem}; use wasmtime::component::wasm_wave::{untyped::UntypedFuncCall, wasm::WasmFunc}; use wasmtime::component::{Component, HasSelf, Linker, ResourceTable, Val}; use wasmtime::*; @@ -112,11 +113,10 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> { let untyped_call = UntypedFuncCall::parse(invoke)?; let exports = collect_export_funcs(&engine, &component); //println!("Exported funcs: {exports:?}"); - let mut find_export = exports.into_iter().filter_map(|(names, func)| { - let func_name = names.last().unwrap(); - (func_name == untyped_call.name()).then_some((names, func)) - }); - let (names, func_type) = &find_export.next().unwrap(); + let (names, func_type) = exports + .into_iter() + .find(|(names, _)| names.last().is_some_and(|name| name == untyped_call.name())) + .with_context(|| format!("no exported function named `{}`", untyped_call.name()))?; let export = names .iter() .fold(None, |instance, name| { @@ -125,20 +125,18 @@ pub fn run(args: RunArgs) -> anyhow::Result<()> { .unwrap(); let instance = linker.instantiate(&mut store, &component)?; - let param_types = WasmFunc::params(func_type).collect::>(); + let param_types = WasmFunc::params(&func_type).collect::>(); let params = untyped_call.to_wasm_params(¶m_types)?; let func = instance.get_func(&mut store, export).unwrap(); let mut results = vec![Val::Bool(false); func_type.results().len()]; - match func.call(&mut store, ¶ms, &mut results) { - Ok(_) => Ok(()), - Err(e) => { - if store.data().exit_called { - Ok(()) - } else { - Err(e) - } + // A replayed `wasi:cli/exit` traps the guest on purpose. + func.call(&mut store, ¶ms, &mut results).or_else(|e| { + if store.data().exit_called { + Ok(()) + } else { + Err(e) } - }?; + })?; if args.trace.is_none() { let trace = store.data().logger.dump_trace(); std::fs::write("trace.out", &trace)?; @@ -154,25 +152,18 @@ fn collect_exports( item: CItem, basename: Vec, ) -> Vec<(Vec, CItem)> { - match item { - CItem::Component(c) => c - .exports(engine) - .flat_map(move |(name, item)| { - let mut names = basename.clone(); - names.push(name.to_string()); - collect_exports(engine, item.ty, names) - }) - .collect(), - CItem::ComponentInstance(c) => c - .exports(engine) - .flat_map(move |(name, item)| { - let mut names = basename.clone(); - names.push(name.to_string()); - collect_exports(engine, item.ty, names) - }) - .collect(), - _ => vec![(basename, item)], - } + let exports: Box)> + '_> = match &item { + CItem::Component(c) => Box::new(c.exports(engine)), + CItem::ComponentInstance(c) => Box::new(c.exports(engine)), + _ => return vec![(basename, item)], + }; + exports + .flat_map(|(name, export)| { + let mut names = basename.clone(); + names.push(name.to_string()); + collect_exports(engine, export.ty, names) + }) + .collect() } fn collect_export_funcs( engine: &Engine, @@ -209,49 +200,35 @@ mod dialog_bindings { }); } +/// Forwards `dialog` host functions of the form `fn(&mut self, dep: u32) -> String`. +macro_rules! forward_dialog_reads { + ($($name:ident),* $(,)?) => { + $( + fn $name(&mut self, dep: u32) -> String { + dialog::$name(dep) + } + )* + }; +} impl dialog_bindings::proxy::util::dialog::Host for crate::run::State { fn print(&mut self, dep: u32, message: String) { dialog::print(dep, &message); } - fn read_string(&mut self, dep: u32) -> String { - dialog::read_string(dep) - } - fn read_u8(&mut self, dep: u32) -> String { - dialog::read_u8(dep) - } - fn read_u16(&mut self, dep: u32) -> String { - dialog::read_u16(dep) - } - fn read_u32(&mut self, dep: u32) -> String { - dialog::read_u32(dep) - } - fn read_u64(&mut self, dep: u32) -> String { - dialog::read_u64(dep) - } - fn read_s8(&mut self, dep: u32) -> String { - dialog::read_s8(dep) - } - fn read_s16(&mut self, dep: u32) -> String { - dialog::read_s16(dep) - } - fn read_s32(&mut self, dep: u32) -> String { - dialog::read_s32(dep) - } - fn read_s64(&mut self, dep: u32) -> String { - dialog::read_s64(dep) - } - fn read_f32(&mut self, dep: u32) -> String { - dialog::read_f32(dep) - } - fn read_f64(&mut self, dep: u32) -> String { - dialog::read_f64(dep) - } - fn read_bool(&mut self, dep: u32) -> String { - dialog::read_bool(dep) - } - fn read_char(&mut self, dep: u32) -> String { - dialog::read_char(dep) - } + forward_dialog_reads!( + read_string, + read_u8, + read_u16, + read_u32, + read_u64, + read_s8, + read_s16, + read_s32, + read_s64, + read_f32, + read_f64, + read_bool, + read_char, + ); fn read_select(&mut self, dep: u32, prompt: String, items: Vec) -> u32 { dialog::read_select(dep, prompt, items) } diff --git a/src/traits/dialog.rs b/src/traits/dialog.rs index 6847b2b..85b930c 100644 --- a/src/traits/dialog.rs +++ b/src/traits/dialog.rs @@ -76,7 +76,7 @@ impl Trait for DialogTrait { fn struct_trait(&self, module_path: &[String], struct_item: &ItemStruct) -> Vec { let mut res = Vec::new(); let struct_name = make_path(module_path, &struct_item.ident.to_string()); - let struct_wit_name = wit_func_name(module_path, &None, &struct_item.ident, &None); + let struct_wit_name = wit_func_name(module_path, None, &struct_item.ident, None); let (impl_generics, ty_generics, where_clause) = struct_item.generics.split_for_impl(); let (field_names, tys, field_wit_names) = match &struct_item.fields { syn::Fields::Unit => (Vec::new(), Vec::new(), Vec::new()), @@ -135,7 +135,7 @@ impl Trait for DialogTrait { syn::Fields::Named(_) => unreachable!(), } }); - let enum_wit_name = wit_func_name(module_path, &None, &enum_item.ident, &None); + let enum_wit_name = wit_func_name(module_path, None, &enum_item.ident, None); res.push(parse_quote! { impl #impl_generics Dialog for #enum_name #ty_generics #where_clause { fn read_value(dep: u32) -> Self { @@ -154,7 +154,7 @@ impl Trait for DialogTrait { fn flag_trait(&self, module_path: &[String], item: &crate::codegen::ItemFlag) -> Vec { let mut res = Vec::new(); let flag_path = make_path(module_path, &item.name.to_string()); - let flag_wit_name = wit_func_name(module_path, &None, &item.name, &None); + let flag_wit_name = wit_func_name(module_path, None, &item.name, None); let flags = &item.flags; let flag_names: Vec<_> = flags .iter() diff --git a/src/traits/wit.rs b/src/traits/wit.rs index 317dc62..dd72d0e 100644 --- a/src/traits/wit.rs +++ b/src/traits/wit.rs @@ -41,7 +41,7 @@ impl Trait for WitTrait { fn struct_trait(&self, module_path: &[String], struct_item: &ItemStruct) -> Vec { let mut res = Vec::new(); let struct_name = make_path(module_path, &struct_item.ident.to_string()); - let struct_wit_name = wit_func_name(module_path, &None, &struct_item.ident, &None); + let struct_wit_name = wit_func_name(module_path, None, &struct_item.ident, None); let (impl_generics, ty_generics, where_clause) = struct_item.generics.split_for_impl(); res.push(parse_quote! { impl #impl_generics WitName for #struct_name #ty_generics #where_clause { @@ -55,7 +55,7 @@ impl Trait for WitTrait { fn enum_trait(&self, module_path: &[String], enum_item: &ItemEnum) -> Vec { let mut res = Vec::new(); let enum_name = make_path(module_path, &enum_item.ident.to_string()); - let enum_wit_name = wit_func_name(module_path, &None, &enum_item.ident, &None); + let enum_wit_name = wit_func_name(module_path, None, &enum_item.ident, None); let (impl_generics, ty_generics, where_clause) = enum_item.generics.split_for_impl(); res.push(parse_quote! { impl #impl_generics WitName for #enum_name #ty_generics #where_clause { @@ -69,7 +69,7 @@ impl Trait for WitTrait { fn flag_trait(&self, module_path: &[String], item: &crate::codegen::ItemFlag) -> Vec { let mut res = Vec::new(); let flag_path = make_path(module_path, &item.name.to_string()); - let flag_wit_name = wit_func_name(module_path, &None, &item.name, &None); + let flag_wit_name = wit_func_name(module_path, None, &item.name, None); res.push(parse_quote! { impl WitName for #flag_path { fn name() -> String {