From ebd0dff445d79626997c39755d8a5d8d5b8c4bdd Mon Sep 17 00:00:00 2001 From: Joel Dice Date: Thu, 25 Jan 2024 15:26:57 -0700 Subject: [PATCH] add `asyncify` option to Rust generator Signed-off-by: Joel Dice --- Cargo.lock | 1 + crates/core/Cargo.toml | 1 + crates/core/src/asyncify.rs | 220 +++++++++++++++++++++++++++++++++++ crates/core/src/lib.rs | 2 + crates/rust-macro/src/lib.rs | 20 ++++ crates/rust/src/interface.rs | 176 +++++++++++++++++++++++++--- crates/rust/src/lib.rs | 13 ++- src/bin/wit-bindgen.rs | 25 ++-- 8 files changed, 432 insertions(+), 26 deletions(-) create mode 100644 crates/core/src/asyncify.rs diff --git a/Cargo.lock b/Cargo.lock index eef2ec9d2..5373215d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2350,6 +2350,7 @@ name = "wit-bindgen-core" version = "0.16.0" dependencies = [ "anyhow", + "indexmap", "wit-component", "wit-parser", ] diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 2ce251748..6076721dc 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -18,3 +18,4 @@ doctest = false wit-parser = { workspace = true } anyhow = { workspace = true } wit-component = { workspace = true } +indexmap = { workspace = true } diff --git a/crates/core/src/asyncify.rs b/crates/core/src/asyncify.rs new file mode 100644 index 000000000..1cd9e80df --- /dev/null +++ b/crates/core/src/asyncify.rs @@ -0,0 +1,220 @@ +use { + indexmap::IndexMap, + std::{collections::HashMap, iter, path::Path}, + wit_parser::{ + Docs, Function, FunctionKind, Interface, InterfaceId, Resolve, Result_, Results, Type, + TypeDef, TypeDefKind, TypeId, TypeOwner, UnresolvedPackage, World, WorldId, WorldItem, + WorldKey, + }, +}; + +struct Asyncify<'a> { + old_resolve: &'a Resolve, + new_resolve: Resolve, + pending: TypeId, + ready: TypeId, + interfaces: HashMap, + functions: HashMap, +} + +impl<'a> Asyncify<'a> { + fn asyncify_world_item( + &mut self, + key: &WorldKey, + item: &WorldItem, + ) -> Vec<(WorldKey, WorldItem)> { + let mut new_key = || match key { + WorldKey::Name(name) => WorldKey::Name(name.clone()), + WorldKey::Interface(old) => { + WorldKey::Interface(if let Some(new) = self.interfaces.get(old).copied() { + new + } else { + let new = self.asyncify_interface(*old); + self.interfaces.insert(*old, new); + new + }) + } + }; + + match item { + WorldItem::Interface(old) => { + vec![( + new_key(), + WorldItem::Interface(if let Some(new) = self.interfaces.get(old).copied() { + new + } else { + let new = self.asyncify_interface(*old); + self.interfaces.insert(*old, new); + new + }), + )] + } + WorldItem::Function(old) => { + let new_key = |suffix| match key { + WorldKey::Name(name) => WorldKey::Name(format!("{name}{suffix}")), + WorldKey::Interface(_) => unreachable!(), + }; + + let (a, b) = if let Some(new) = self.functions.get(key).cloned() { + new + } else { + let new = self.asyncify_function(old); + self.functions.insert(key.clone(), new.clone()); + new + }; + vec![ + (new_key("-isyswasfa"), WorldItem::Function(a)), + (new_key("-isyswasfa-result"), WorldItem::Function(b)), + ] + } + WorldItem::Type(old) => vec![(new_key(), WorldItem::Type(*old))], + } + } + + fn asyncify_interface(&mut self, interface: InterfaceId) -> InterfaceId { + let old = &self.old_resolve.interfaces[interface]; + let functions = old + .functions + .iter() + .flat_map(|(_, function)| { + let (a, b) = self.asyncify_function(function); + [(a.name.clone(), a), (b.name.clone(), b)] + }) + .collect(); + + self.new_resolve.interfaces.alloc(Interface { + name: old.name.as_ref().map(|s| format!("{s}-isyswasfa")), + types: old.types.clone(), + functions, + docs: old.docs.clone(), + package: old.package, + }) + } + + fn asyncify_function(&mut self, function: &Function) -> (Function, Function) { + ( + Function { + name: format!("{}-isysasfa", function.name), + kind: function.kind.clone(), + params: function.params.clone(), + results: match &function.results { + Results::Anon(ty) => { + Results::Anon(Type::Id(self.new_resolve.types.alloc(TypeDef { + name: None, + kind: TypeDefKind::Result(Result_ { + ok: Some(*ty), + err: Some(Type::Id(self.pending)), + }), + owner: TypeOwner::None, + docs: Docs::default(), + }))) + } + Results::Named(_) => { + todo!("handle functions returning multiple named results") + } + }, + docs: function.docs.clone(), + }, + Function { + name: format!("{}-isysasfa-result", function.name), + kind: function.kind.clone(), + params: vec![("pending".into(), Type::Id(self.ready))], + results: function.results.clone(), + docs: function.docs.clone(), + }, + ) + } +} + +pub fn asyncify(resolve: &Resolve, world: WorldId, poll_suffix: &str) -> (Resolve, WorldId) { + let old_world = &resolve.worlds[world]; + + let mut new_resolve = resolve.clone(); + + let isyswasfa_package = new_resolve + .push( + UnresolvedPackage::parse( + &Path::new("isyswasfa.wit"), + include_str!("../../../../wit/isyswasfa.wit"), + ) + .unwrap(), + ) + .unwrap(); + + let isyswasfa_interface = new_resolve.packages[isyswasfa_package].interfaces["isyswasfa"]; + + let poll_input = new_resolve.interfaces[isyswasfa_interface].types["poll-input"]; + let list_poll_input = new_resolve.types.alloc(TypeDef { + name: None, + kind: TypeDefKind::List(Type::Id(poll_input)), + owner: TypeOwner::None, + docs: Docs::default(), + }); + + let poll_output = new_resolve.interfaces[isyswasfa_interface].types["poll-output"]; + let list_poll_output = new_resolve.types.alloc(TypeDef { + name: None, + kind: TypeDefKind::List(Type::Id(poll_output)), + owner: TypeOwner::None, + docs: Docs::default(), + }); + + let new_world = new_resolve.worlds.alloc(World { + name: format!("{}-isyswasfa", old_world.name), + imports: IndexMap::new(), + exports: IndexMap::new(), + package: old_world.package, + docs: old_world.docs.clone(), + includes: Vec::new(), + include_names: Vec::new(), + }); + + let poll_function_name = format!("isyswasfa-poll-{poll_suffix}"); + let poll_function = Function { + name: poll_function_name.clone(), + kind: FunctionKind::Freestanding, + params: vec![("input".to_owned(), Type::Id(list_poll_input))], + results: Results::Anon(Type::Id(list_poll_output)), + docs: Docs::default(), + }; + + let pending = new_resolve.interfaces[isyswasfa_interface].types["pending"]; + let ready = new_resolve.interfaces[isyswasfa_interface].types["ready"]; + + let mut asyncify = Asyncify { + old_resolve: resolve, + new_resolve, + pending, + ready, + interfaces: HashMap::new(), + functions: HashMap::new(), + }; + + let imports = old_world + .imports + .iter() + .flat_map(|(key, item)| asyncify.asyncify_world_item(key, item)) + .chain(iter::once(( + WorldKey::Interface(isyswasfa_interface), + WorldItem::Interface(isyswasfa_interface), + ))) + .collect(); + + let exports = old_world + .imports + .iter() + .flat_map(|(key, item)| asyncify.asyncify_world_item(key, item)) + .chain(iter::once(( + WorldKey::Name(poll_function_name), + WorldItem::Function(poll_function), + ))) + .collect(); + + { + let new_world = &mut asyncify.new_resolve.worlds[new_world]; + new_world.imports = imports; + new_world.exports = exports; + } + + (asyncify.new_resolve, new_world) +} diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 5cbe900ac..663e0cc78 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -9,6 +9,8 @@ mod ns; pub use ns::Ns; pub mod source; pub use source::{Files, Source}; +pub mod asyncify; +pub use asyncify::asyncify; #[derive(Default, Copy, Clone, PartialEq, Eq, Debug)] pub enum Direction { diff --git a/crates/rust-macro/src/lib.rs b/crates/rust-macro/src/lib.rs index a2dfe4e42..6ad0df359 100644 --- a/crates/rust-macro/src/lib.rs +++ b/crates/rust-macro/src/lib.rs @@ -94,6 +94,15 @@ impl Parse for Config { .collect() } Opt::With(with) => opts.with.extend(with), + Opt::Asyncify(suffix) => { + if opts.asyncify.is_some() { + return Err(Error::new( + suffix.span(), + "cannot specify second asyncify suffix", + )); + } + opts.asyncify = Some(suffix.value()) + } } } } else { @@ -107,6 +116,11 @@ impl Parse for Config { let world = resolve .select_world(pkg, world.as_deref()) .map_err(|e| Error::new(call_site, format!("{e:?}")))?; + let (resolve, world) = if let Some(suffix) = opts.asyncify.as_deref() { + wit_bindgen_core::asyncify(&resolve, world, suffix) + } else { + (resolve, world) + }; Ok(Config { opts, resolve, @@ -186,6 +200,7 @@ mod kw { syn::custom_keyword!(export_prefix); syn::custom_keyword!(additional_derives); syn::custom_keyword!(with); + syn::custom_keyword!(asyncify); } #[derive(Clone)] @@ -246,6 +261,7 @@ enum Opt { // Parse as paths so we can take the concrete types/macro names rather than raw strings AdditionalDerives(Vec), With(HashMap), + Asyncify(syn::LitStr), } impl Parse for Opt { @@ -336,6 +352,10 @@ impl Parse for Opt { input.parse::()?; input.parse::()?; Ok(Opt::ExportPrefix(input.parse()?)) + } else if l.peek(kw::asyncify) { + input.parse::()?; + input.parse::()?; + Ok(Opt::Asyncify(input.parse()?)) } else if l.peek(kw::additional_derives) { input.parse::()?; input.parse::()?; diff --git a/crates/rust/src/interface.rs b/crates/rust/src/interface.rs index 59f8891b9..fde8b3590 100644 --- a/crates/rust/src/interface.rs +++ b/crates/rust/src/interface.rs @@ -80,6 +80,7 @@ impl InterfaceGenerator<'_> { funcs: impl Iterator + Clone, ) -> Result<()> { let mut traits = BTreeMap::new(); + let mut asyncified_traits = BTreeMap::new(); for func in funcs { if self.gen.skip.contains(&func.name) { @@ -92,30 +93,34 @@ impl InterfaceGenerator<'_> { self.generate_guest_export(func); self.src.push_str("};\n"); - // Next generate a trait signature for this method and insert it - // into `traits`. Note that `traits` will have a trait-per-resource. - let (trait_name, local_impl_name, export_key) = match func.kind { + let info = |me: &mut Self, prefix| match func.kind { FunctionKind::Freestanding => ( - "Guest".to_string(), - "_GuestImpl".to_string(), - self.export_key(None), + format!("{prefix}Guest"), + format!("_{prefix}GuestImpl"), + me.export_key(None), ), FunctionKind::Method(id) | FunctionKind::Constructor(id) | FunctionKind::Static(id) => { - let resource_name = self.resolve.types[id].name.as_deref().unwrap(); + let resource_name = me.resolve.types[id].name.as_deref().unwrap(); let camel = resource_name.to_upper_camel_case(); - let trait_name = format!("Guest{camel}"); - let export_key = self.export_key(Some(&resource_name)); - let local_impl_name = format!("_{camel}Impl"); + let trait_name = format!("{prefix}Guest{camel}"); + let export_key = me.export_key(Some(&resource_name)); + let local_impl_name = format!("_{prefix}{camel}Impl"); (trait_name, local_impl_name, export_key) } }; - let (_, _, methods) = - traits - .entry(export_key) - .or_insert((trait_name, local_impl_name, Vec::new())); + // Next generate a trait signature for this method and insert it + // into `traits`. Note that `traits` will have a trait-per-resource. + let (trait_name, local_impl_name, export_key) = info(self, ""); + + let (_, _, methods, impls) = traits.entry(export_key).or_insert(( + trait_name, + local_impl_name, + Vec::new(), + Vec::new(), + )); let prev = mem::take(&mut self.src); let mut sig = FnSig { use_item_name: true, @@ -130,6 +135,71 @@ impl InterfaceGenerator<'_> { self.src.push_str(";\n"); let trait_method = mem::replace(&mut self.src, prev); methods.push(trait_method); + + if let Some(suffix) = self.gen.opts.asyncify.clone() { + let prev = mem::take(&mut self.src); + self.print_signature(func, TypeMode::Owned, &sig); + if func.name == format!("isyswasfa-poll{suffix}") { + self.src.push_str("{ isyswasfa_guest::poll(input) }\n"); + } else if let Some(prefix) = func.name.strip_suffix("-isyswasfa") { + let (trait_name, local_impl_name, export_key) = info(self, "Isyswasfa"); + + { + let (_, _, methods) = asyncified_traits.entry(export_key).or_insert(( + trait_name.clone(), + local_impl_name.clone(), + Vec::new(), + )); + + let sig = FnSig { + async_: true, + ..sig.clone() + }; + let func = &Function { + name: prefix.into(), + kind: func.kind.clone(), + params: func.params.clone(), + results: if let Results::Anon(Type::Id(id)) = &func.results { + if let TypeDefKind::Result(Result_ { ok: Some(ok), .. }) = + &self.resolve.types[*id].kind + { + Results::Anon(*ok) + } else { + unreachable!() + } + } else { + unreachable!() + }, + docs: func.docs.clone(), + }; + let prev = mem::take(&mut self.src); + self.print_signature(func, TypeMode::Owned, &sig); + self.src.push_str(";\n"); + + let trait_method = mem::replace(&mut self.src, prev); + methods.push(trait_method); + } + + let params = func + .params + .iter() + .map(|(name, _)| to_rust_ident(name)) + .collect::>() + .join(", "); + + uwriteln!( + self.src, + "{{ isyswasfa_guest::first_poll(<{local_impl_name} as {trait_name}>::{}({params})) }}", + to_rust_ident(prefix) + ); + } else if func.name.ends_with("-isyswasfa-result") { + self.src.push_str("{ isyswasfa_guest::get_ready(ready)) }"); + } else { + unreachable!() + } + let impl_ = mem::replace(&mut self.src, prev); + impls.push(impl_); + } } // Once all the traits have been assembled then they can be emitted. @@ -137,7 +207,31 @@ impl InterfaceGenerator<'_> { // Additionally alias the user-configured item for each trait here as // there's only one implementation of this trait and it must be // pre-configured. - for (export_key, (trait_name, local_impl_name, methods)) in traits { + for (export_key, (trait_name, local_impl_name, methods, impls)) in traits { + if self.gen.opts.asyncify.is_some() { + uwriteln!(self.src, "struct {local_impl_name};"); + uwriteln!(self.src, "impl {trait_name} for {local_impl_name} {{",); + for impl_ in impls { + self.src.push_str(&impl_); + } + uwriteln!(self.src, "}}"); + } else { + let impl_name = self.gen.lookup_export(&export_key)?; + let path_to_root = self.path_to_root(); + uwriteln!( + self.src, + "use {path_to_root}{impl_name} as {local_impl_name};" + ); + } + + uwriteln!(self.src, "pub trait {trait_name} {{"); + for method in methods { + self.src.push_str(&method); + } + uwriteln!(self.src, "}}"); + } + + for (export_key, (trait_name, local_impl_name, methods)) in asyncified_traits { let impl_name = self.gen.lookup_export(&export_key)?; let path_to_root = self.path_to_root(); uwriteln!( @@ -145,7 +239,7 @@ impl InterfaceGenerator<'_> { "use {path_to_root}{impl_name} as {local_impl_name};" ); - uwriteln!(self.src, "pub trait {trait_name} {{"); + uwriteln!(self.src, "#[async_trait(?Send)]\npub trait {trait_name} {{"); for method in methods { self.src.push_str(&method); } @@ -307,6 +401,55 @@ impl InterfaceGenerator<'_> { self.src.push_str("}\n"); self.src.push_str("}\n"); + if self.gen.opts.asyncify.is_some() { + if let Some(prefix) = func.name.strip_suffix("-isyswasfa") { + sig.async_ = true; + self.src.push_str("#[allow(unused_unsafe, clippy::all)]\n"); + let params = self + .print_signature( + &Function { + name: prefix.into(), + kind: func.kind.clone(), + params: func.params.clone(), + results: if let Results::Anon(Type::Id(id)) = &func.results { + if let TypeDefKind::Result(Result_ { ok: Some(ok), .. }) = + &self.resolve.types[*id].kind + { + Results::Anon(*ok) + } else { + unreachable!() + } + } else { + unreachable!() + }, + docs: func.docs.clone(), + }, + param_mode, + &sig, + ) + .join(", "); + + let isyswasfa = to_rust_ident(&func.name); + let isyswasfa = match func.kind { + FunctionKind::Freestanding => isyswasfa, + FunctionKind::Method(_) + | FunctionKind::Static(_) + | FunctionKind::Constructor(_) => format!("Self::{isyswasfa}"), + }; + let isyswasfa_result = &format!("{isyswasfa}_result"); + + uwrite!( + self.src, + " + match {isyswasfa}({params}) {{ + Ok(result) => result, + Err(pending) => {isyswasfa_result}(isyswasfa_guest::await_ready(pending).await), + }} + " + ); + } + } + match func.kind { FunctionKind::Freestanding => {} FunctionKind::Method(_) | FunctionKind::Static(_) | FunctionKind::Constructor(_) => { @@ -437,7 +580,6 @@ impl InterfaceGenerator<'_> { self.src.push_str("};\n"); } } - pub fn generate_stub( &mut self, resource: Option, diff --git a/crates/rust/src/lib.rs b/crates/rust/src/lib.rs index 8583a02d9..d417cd6f1 100644 --- a/crates/rust/src/lib.rs +++ b/crates/rust/src/lib.rs @@ -178,6 +178,9 @@ pub struct Opts { /// Remapping of interface names to rust module names. #[cfg_attr(feature = "clap", arg(long, value_parser = parse_with, default_value = ""))] pub with: HashMap, + + #[cfg_attr(feature = "clap", arg(long))] + pub asyncify: Option, } impl Opts { @@ -262,6 +265,9 @@ impl RustWasm { } fn lookup_export(&self, key: &ExportKey) -> Result { + if self.opts.asyncify.is_some() { + return Ok("Isyswasfa".to_owned()); + } if let Some(key) = self.opts.exports.get(key) { return Ok(key.clone()); } @@ -509,6 +515,11 @@ impl WorldGenerator for RustWasm { ); if self.opts.stubs { + if self.opts.asyncify.is_some() { + // this shouldn't be hard to implement, but isn't a priority at the moment: + todo!("stubs not yet supported in combination with asyncify"); + } + self.src.push_str("\n#[derive(Debug)]\npub struct Stub;\n"); let world_id = world; let world = &resolve.worlds[world]; @@ -687,7 +698,7 @@ impl fmt::Display for Ownership { } } -#[derive(Default)] +#[derive(Default, Clone)] struct FnSig { async_: bool, unsafe_: bool, diff --git a/src/bin/wit-bindgen.rs b/src/bin/wit-bindgen.rs index 5eb78a1b5..5d7e36877 100644 --- a/src/bin/wit-bindgen.rs +++ b/src/bin/wit-bindgen.rs @@ -91,22 +91,25 @@ struct Common { fn main() -> Result<()> { let mut files = Files::default(); - let (generator, opt) = match Opt::parse() { + let (generator, opt, asyncify) = match Opt::parse() { #[cfg(feature = "markdown")] - Opt::Markdown { opts, args } => (opts.build(), args), + Opt::Markdown { opts, args } => (opts.build(), args, None), #[cfg(feature = "c")] - Opt::C { opts, args } => (opts.build(), args), + Opt::C { opts, args } => (opts.build(), args, None), #[cfg(feature = "rust")] - Opt::Rust { opts, args } => (opts.build(), args), + Opt::Rust { opts, args } => { + let asyncify = opts.asyncify.clone(); + (opts.build(), args, asyncify) + } #[cfg(feature = "teavm-java")] - Opt::TeavmJava { opts, args } => (opts.build(), args), + Opt::TeavmJava { opts, args } => (opts.build(), args, None), #[cfg(feature = "go")] - Opt::TinyGo { opts, args } => (opts.build(), args), + Opt::TinyGo { opts, args } => (opts.build(), args, None), #[cfg(feature = "csharp")] - Opt::CSharp { opts, args } => (opts.build(), args), + Opt::CSharp { opts, args } => (opts.build(), args, None), }; - gen_world(generator, &opt, &mut files)?; + gen_world(generator, &opt, &mut files, asyncify)?; for (name, contents) in files.iter() { let dst = match &opt.out_dir { @@ -153,6 +156,7 @@ fn gen_world( mut generator: Box, opts: &Common, files: &mut Files, + asyncify: Option, ) -> Result<()> { let mut resolve = Resolve::default(); let pkg = if opts.wit.is_dir() { @@ -161,6 +165,11 @@ fn gen_world( resolve.push(UnresolvedPackage::parse_file(&opts.wit)?)? }; let world = resolve.select_world(pkg, opts.world.as_deref())?; + let (resolve, world) = if let Some(suffix) = asyncify.as_deref() { + wit_bindgen_core::asyncify(&resolve, world, suffix) + } else { + (resolve, world) + }; if let Err(e) = generator.generate(&resolve, world, files) { eprintln!( "{e:?}\n\n\