use std::collections::HashMap;
use anyhow::{bail, Context};
use heck::{ToKebabCase, ToUpperCamelCase};
use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, ToTokens, TokenStreamExt};
use syn::{
parse_macro_input, punctuated::Punctuated, visit_mut::VisitMut, ImplItemFn, ItemEnum,
ItemStruct, ItemType, LitStr, PathSegment, ReturnType, Token,
};
use tracing::debug;
use tracing_subscriber::EnvFilter;
use wit_parser::{Resolve, WorldKey};
mod bindgen_visitor;
use bindgen_visitor::WitBindgenOutputVisitor;
mod config;
use config::ProviderBindgenConfig;
mod rust;
mod vendor;
use vendor::wasmtime_component_macro::bindgen::expand as expand_wasmtime_component;
mod wit;
use wit::{
translate_export_fn_for_lattice, WitFunctionName, WitInterfacePath, WitNamespaceName,
WitPackageName,
};
use crate::wit::translate_import_fn_for_lattice;
mod wrpc;
const EXPORTS_MODULE_NAME: &str = "exports";
type ImplStructName = String;
type WasmcloudContract = String;
type LatticeExposedInterface = (WitNamespaceName, WitPackageName, WitFunctionName);
type StructName = String;
type StructLookup = HashMap<StructName, (Punctuated<PathSegment, Token![::]>, ItemStruct)>;
type EnumName = String;
type EnumLookup = HashMap<EnumName, (Punctuated<PathSegment, Token![::]>, ItemEnum)>;
type TypeName = String;
type TypeLookup = HashMap<TypeName, (Punctuated<PathSegment, Token![::]>, ItemType)>;
type WitTraitName = String;
#[derive(Debug, Clone)]
struct ExportedLatticeMethod {
operation_name: LitStr,
func_name: Ident,
invocation_args: Vec<(Ident, TokenStream)>,
invocation_return: ReturnType,
}
#[proc_macro]
pub fn generate(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
let _ = tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env())
.try_init();
let cfg = parse_macro_input!(input as ProviderBindgenConfig);
let contract_ident = LitStr::new(&cfg.contract, Span::call_site());
let wit_bindgen_cfg = cfg
.wit_bindgen_cfg
.as_ref()
.context("configuration to pass to WIT bindgen is missing")
.expect("failed to parse WIT bindgen configuration");
let mut imported_iface_invocation_methods: Vec<TokenStream> = Vec::new();
for (_, world) in wit_bindgen_cfg.resolve.worlds.iter() {
for (import_key, _) in world.imports.iter() {
if let WorldKey::Interface(iface_id) = import_key {
let iface = &wit_bindgen_cfg.resolve.interfaces[*iface_id];
if iface
.package
.map(|p| &wit_bindgen_cfg.resolve.packages[p].name)
.is_some_and(is_ignored_invocation_handler_pkg)
{
continue;
}
for (iface_fn_name, iface_fn) in iface.functions.iter() {
debug!("processing imported interface function: [{iface_fn_name}]");
imported_iface_invocation_methods.push(
translate_import_fn_for_lattice(iface, iface_fn_name, iface_fn, &cfg)
.expect("failed to translate export fn"),
);
}
}
}
}
let bindgen_tokens: TokenStream =
expand_wasmtime_component(wit_bindgen_cfg).unwrap_or_else(syn::Error::into_compile_error);
let mut bindgen_ast: syn::File =
syn::parse2(bindgen_tokens).expect("failed to parse wit-bindgen generated code as file");
let mut visitor = WitBindgenOutputVisitor::new(&cfg);
visitor.visit_file_mut(&mut bindgen_ast);
let methods_by_iface = build_lattice_methods_by_wit_interface(
&visitor.serde_extended_structs,
&visitor.type_lookup,
&visitor.export_trait_methods,
&cfg,
)
.expect("failed to build lattice methods from WIT interfaces");
let impl_struct_name = Ident::new_raw(cfg.impl_struct.as_str(), Span::call_site());
let mut interface_dispatch_wrpc_match_arms: Vec<TokenStream> = Vec::new();
let mut iface_tokens = TokenStream::new();
for (wit_iface_name, methods) in methods_by_iface.iter() {
let wit_iface = Ident::new(wit_iface_name, Span::call_site());
let operation_names = methods
.clone()
.into_iter()
.map(|lm| lm.operation_name)
.collect::<Vec<LitStr>>();
let func_names = methods
.clone()
.into_iter()
.map(|lm| lm.func_name)
.collect::<Vec<Ident>>();
let invocation_args_with_types = methods
.clone()
.into_iter()
.map(|lm| {
let arg_tokens = lm
.invocation_args
.iter()
.map(|(ident, ty)| quote!(#ident: #ty))
.collect::<Vec<TokenStream>>();
quote::quote!(#( #arg_tokens ),*)
})
.collect::<Vec<TokenStream>>();
let invocation_returns = methods
.clone()
.into_iter()
.map(|lm| lm.invocation_return)
.collect::<Vec<ReturnType>>();
iface_tokens.append_all(quote!(
#[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
pub trait #wit_iface {
fn contract_id() -> &'static str {
#contract_ident
}
#(
async fn #func_names (
&self,
ctx: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Context,
#invocation_args_with_types
) #invocation_returns;
)*
}
));
let (wrpc_input_parsing_statements, post_self_args, result_encode_tokens) = methods.clone().into_iter().fold(
(Vec::<TokenStream>::new(), Vec::<TokenStream>::new(), Vec::<TokenStream>::new()),
|mut acc, lm| {
let mut input_decoding_lines = Vec::<TokenStream>::new();
for (arg_name, arg_type) in lm.invocation_args.iter() {
let arg_name_lit = LitStr::new(&arg_name.to_string(), Span::call_site());
let arg_ty = arg_type.to_token_stream();
input_decoding_lines.push(quote::quote!(
let mut #arg_name = ::wasmcloud_provider_wit_bindgen::deps::bytes::BytesMut::new();
params
.pop()
.ok_or_else(|| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("missing expected parameter [{}]", #arg_name_lit)))?
.encode(&mut #arg_name)
.await
.map_err(|e| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("failed to encode parameter [{}]: {e}", #arg_name_lit)))?;
let (#arg_name, _): (#arg_ty, _) = ::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::Receive::receive::<::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::DemuxStream>(#arg_name, &mut ::wasmcloud_provider_wit_bindgen::deps::futures::stream::empty(), None)
.await
.map_err(|e| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("failed to receive parameter [{}]: {e}", #arg_name_lit)))?;
));
}
acc.0.push(quote::quote!(#( #input_decoding_lines );*));
let arg_idents = vec![Ident::new("ctx", Span::call_site())]
.into_iter()
.chain(lm.invocation_args.iter().map(|(name, _)| name.clone()))
.collect::<Vec<Ident>>();
acc.1.push(quote!(#( #arg_idents ),*));
acc.2.push(match lm.invocation_return {
syn::ReturnType::Type(_, _) => {
quote!(result
.encode(&mut res)
.await
.map_err(|e| {
::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(
format!("failed to encode result of operation [{operation}]: {e}")
)
})?)
}
syn::ReturnType::Default => {
quote!(result
.encode(&mut res)
.await
.map_err(|e| {
::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(
format!("failed to encode result of operation [{operation}]: {e}")
)
})?)
},
});
acc
},
);
interface_dispatch_wrpc_match_arms.push(quote!(
#(
operation @ #operation_names => {
#wrpc_input_parsing_statements
let result = #wit_iface::#func_names(
self,
#post_self_args
)
.await;
let mut res = ::wasmcloud_provider_wit_bindgen::deps::bytes::BytesMut::new();
#result_encode_tokens;
Ok(res.to_vec())
}
)*
));
}
let types: Vec<TokenStream> = visitor
.type_lookup
.iter()
.filter_map(|(_, (_, ty))| {
if visitor
.serde_extended_structs
.contains_key(&ty.ident.to_string())
|| visitor
.serde_extended_enums
.contains_key(&ty.ident.to_string())
{
None
} else {
Some(ty.to_token_stream())
}
})
.collect();
let structs: Vec<TokenStream> = visitor
.serde_extended_structs
.iter()
.map(|(_, (_, s))| s.to_token_stream())
.collect();
let enums: Vec<TokenStream> = visitor
.serde_extended_enums
.iter()
.map(|(_, (_, s))| s.to_token_stream())
.collect();
let wrpc_impl_tokens = build_wrpc_impls(&impl_struct_name, &wit_bindgen_cfg.resolve)
.expect("failed to build provider-sdk wrpc implementation");
let tokens = quote!(
#iface_tokens
#(
#types
)*
#(
#structs
)*
#(
#enums
)*
#[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
trait WasmcloudCapabilityProvider {
async fn receive_link_config_as_source(
&self,
link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
Ok(())
}
async fn receive_link_config_as_target(
&self,
link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
Ok(())
}
async fn delete_link(
&self,
actor_id: &str
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
Ok(())
}
async fn shutdown(&self) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
Ok(())
}
}
#[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderHandler for #impl_struct_name {
async fn receive_link_config_as_source(
&self,
link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
WasmcloudCapabilityProvider::receive_link_config_as_source(self, link_config).await
}
async fn receive_link_config_as_target(
&self,
link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
WasmcloudCapabilityProvider::receive_link_config_as_target(self, link_config).await
}
async fn delete_link(
&self,
actor_id: &str
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
WasmcloudCapabilityProvider::delete_link(self, actor_id).await
}
async fn shutdown(&self) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
WasmcloudCapabilityProvider::shutdown(self).await
}
}
impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Provider for #impl_struct_name {}
pub struct InvocationHandler {
wrpc_client: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::core::wrpc::Client
}
impl InvocationHandler {
pub fn new(target: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::core::ComponentId) -> Self {
let connection = ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::get_connection();
Self { wrpc_client: connection.get_wrpc_client(&target) }
}
#(
#imported_iface_invocation_methods
)*
}
#wrpc_impl_tokens
#[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::WrpcDispatch for #impl_struct_name {
async fn dispatch_wrpc_dynamic<'a>(
&'a self,
ctx: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Context,
operation: String,
mut params: Vec<::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::Value>,
) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationResult<Vec<u8>> {
use ::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::{Encode, Receive};
use ::wasmcloud_provider_wit_bindgen::deps::anyhow::Context as _;
match operation.as_str() {
#(
#interface_dispatch_wrpc_match_arms
)*
_ => Err(::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Malformed(format!(
"Invalid operation name [{operation}]"
)).into())
}
}
}
);
tokens.into()
}
fn build_lattice_methods_by_wit_interface(
struct_lookup: &StructLookup,
type_lookup: &TypeLookup,
export_trait_methods: &HashMap<WitInterfacePath, Vec<ImplItemFn>>,
bindgen_cfg: &ProviderBindgenConfig,
) -> anyhow::Result<HashMap<WitTraitName, Vec<ExportedLatticeMethod>>> {
let mut methods_by_name: HashMap<WitInterfacePath, Vec<ExportedLatticeMethod>> = HashMap::new();
for (wit_iface_name, funcs) in export_trait_methods.iter() {
for trait_method in funcs.iter() {
let wit_operation = match wit_iface_name.split('.').collect::<Vec<&str>>()[..] {
[wit_ns, wit_pkg, iface] => {
format!(
"{}:{}/{}.{}",
wit_ns.to_kebab_case(),
wit_pkg.to_kebab_case(),
iface.to_kebab_case(),
trait_method.sig.ident.to_string().to_kebab_case()
)
}
_ => bail!("unexpected interface path, expected 3 components"),
};
let operation_name = LitStr::new(&wit_operation, trait_method.sig.ident.span());
let lattice_method = translate_export_fn_for_lattice(
bindgen_cfg,
operation_name,
trait_method,
struct_lookup,
type_lookup,
)?;
let wit_iface_upper_camel = wit_iface_name
.split('.')
.map(|v| v.to_upper_camel_case())
.collect::<String>();
methods_by_name
.entry(wit_iface_upper_camel)
.or_default()
.push(lattice_method);
}
}
Ok(methods_by_name)
}
fn is_ignored_invocation_handler_pkg(pkg: &wit_parser::PackageName) -> bool {
matches!(
(pkg.namespace.as_ref(), pkg.name.as_ref()),
("wasmcloud", "bus") | ("wasi", "io")
)
}
fn build_wrpc_impls(impl_struct_name: &Ident, resolve: &Resolve) -> anyhow::Result<TokenStream> {
let mapping = crate::wrpc::generate_wrpc_nats_subject_to_fn_mapping(resolve)
.context("failed to generate wrpc NATS subject mappings")?;
let mut insertion_lines: Vec<TokenStream> = Vec::new();
for crate::wrpc::WrpcExport {
wit_ns,
wit_pkg,
wit_iface,
wit_iface_fn,
types,
} in mapping.into_iter()
{
let wit_ns = LitStr::new(&wit_ns, Span::call_site());
let wit_pkg = LitStr::new(&wit_pkg, Span::call_site());
let wit_iface = LitStr::new(&wit_iface, Span::call_site());
let wit_iface_fn = LitStr::new(&wit_iface_fn, Span::call_site());
let world_key_name = LitStr::new(&types.0, Span::call_site());
let function_name = LitStr::new(&types.1, Span::call_site());
let dynamic_fn = LitStr::new(
&serde_json::to_string::<wrpc_types::DynamicFunction>(&types.2).context("failed to deserialize dynamic function with world_key_name [{world_key_name}], function name [{function_name}]")?,
Span::call_site(),
);
insertion_lines.push(quote!(
mapping.insert(
format!("{lattice_name}.{component_id}.wrpc.{wrpc_version}.{}:{}/{}.{}", #wit_ns, #wit_pkg, #wit_iface, #wit_iface_fn),
(#world_key_name.into(), #function_name.into(), ::wasmcloud_provider_wit_bindgen::deps::serde_json::from_slice::<::wasmcloud_provider_wit_bindgen::deps::wrpc_types::DynamicFunction>(#dynamic_fn.as_bytes()).expect("failed to deserialize DynamicFunction")),
);
));
}
let tokens = quote!(
use ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::{WrpcNats, WrpcNatsSubject, WorldKeyName, WitFunction};
use ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::ProviderInitResult;
#[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
impl WrpcNats for #impl_struct_name {
async fn incoming_wrpc_invocations_by_subject(
&self,
lattice_name: impl AsRef<str> + Send,
component_id: impl AsRef<str> + Send,
wrpc_version: impl AsRef<str> + Send,
) -> ProviderInitResult<
::std::collections::HashMap<WrpcNatsSubject, (WorldKeyName, WitFunction, ::wasmcloud_provider_wit_bindgen::deps::wrpc_types::DynamicFunction)>
> {
let lattice_name = lattice_name.as_ref();
let wrpc_version = wrpc_version.as_ref();
let component_id = component_id.as_ref();
let mut mapping = ::std::collections::HashMap::new();
#(
#insertion_lines
)*
Ok(mapping)
}
}
);
Ok(tokens)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use anyhow::Context;
use proc_macro2::{Span, TokenTree};
use quote::quote;
use syn::{parse_quote, ImplItemFn, LitStr, ReturnType};
use crate::wit::{extract_witified_map, translate_export_fn_for_lattice};
use crate::ProviderBindgenConfig;
#[test]
fn parse_witified_map_type() -> anyhow::Result<()> {
extract_witified_map(
"e!(Vec<(String, String)>)
.into_iter()
.collect::<Vec<TokenTree>>(),
)
.context("failed to parse WIT-ified map type Vec<(String, String)>")?;
Ok(())
}
#[test]
fn parse_witified_map_in_fn() -> anyhow::Result<()> {
let trait_fn: ImplItemFn = parse_quote!(
fn baz(test_map: Vec<(String, String)>) {}
);
let bindgen_cfg = ProviderBindgenConfig {
impl_struct: "None".into(),
contract: "wasmcloud:test".into(),
wit_ns: Some("test".into()),
wit_pkg: Some("foo".into()),
exposed_interface_allow_list: Default::default(),
exposed_interface_deny_list: Default::default(),
wit_bindgen_cfg: None, replace_witified_maps: true,
};
let operation_name = LitStr::new("wasmcloud:test/test.foo", Span::call_site());
let lm = translate_export_fn_for_lattice(
&bindgen_cfg,
operation_name.clone(),
&trait_fn,
&HashMap::new(), &HashMap::new(), )?;
assert_eq!(lm.operation_name, operation_name);
assert_eq!(lm.invocation_args.len(), 1);
assert_eq!(
lm.invocation_return,
syn::parse2::<ReturnType>(quote::quote!())?
);
Ok(())
}
}