1use std::collections::HashMap;
30
31use anyhow::{bail, Context};
32use heck::{ToKebabCase, ToUpperCamelCase};
33use proc_macro2::{Ident, Span, TokenStream};
34use quote::{quote, ToTokens, TokenStreamExt};
35use syn::{
36 parse_macro_input, punctuated::Punctuated, visit_mut::VisitMut, ImplItemFn, ItemEnum,
37 ItemStruct, ItemType, LitStr, PathSegment, ReturnType, Token,
38};
39use tracing::debug;
40use tracing_subscriber::EnvFilter;
41use wit_parser::{Resolve, WorldKey};
42
43mod bindgen_visitor;
44use bindgen_visitor::WitBindgenOutputVisitor;
45
46mod config;
47use config::ProviderBindgenConfig;
48
49mod rust;
50
51mod vendor;
52use vendor::wasmtime_component_macro::bindgen::expand as expand_wasmtime_component;
53
54mod wit;
55use wit::{
56 translate_export_fn_for_lattice, WitFunctionName, WitInterfacePath, WitNamespaceName,
57 WitPackageName,
58};
59
60use crate::wit::translate_import_fn_for_lattice;
61
62mod wrpc;
63
64const EXPORTS_MODULE_NAME: &str = "exports";
66
67type ImplStructName = String;
68type WasmcloudContract = String;
69
70type LatticeExposedInterface = (WitNamespaceName, WitPackageName, WitFunctionName);
72
73type StructName = String;
74type StructLookup = HashMap<StructName, (Punctuated<PathSegment, Token![::]>, ItemStruct)>;
75
76type EnumName = String;
77type EnumLookup = HashMap<EnumName, (Punctuated<PathSegment, Token![::]>, ItemEnum)>;
78
79type TypeName = String;
80type TypeLookup = HashMap<TypeName, (Punctuated<PathSegment, Token![::]>, ItemType)>;
81
82type WitTraitName = String;
84
85#[derive(Debug, Clone)]
90struct ExportedLatticeMethod {
91 operation_name: LitStr,
93
94 func_name: Ident,
96
97 invocation_args: Vec<(Ident, TokenStream)>,
99
100 invocation_return: ReturnType,
102}
103
104#[proc_macro]
106pub fn generate(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
107 let _ = tracing_subscriber::fmt()
108 .with_env_filter(EnvFilter::from_default_env())
109 .try_init();
110
111 let cfg = parse_macro_input!(input as ProviderBindgenConfig);
113 let contract_ident = LitStr::new(&cfg.contract, Span::call_site());
114
115 let wit_bindgen_cfg = cfg
118 .wit_bindgen_cfg
119 .as_ref()
120 .context("configuration to pass to WIT bindgen is missing")
121 .expect("failed to parse WIT bindgen configuration");
122
123 let mut imported_iface_invocation_methods: Vec<TokenStream> = Vec::new();
126 for (_, world) in wit_bindgen_cfg.resolve.worlds.iter() {
127 for (import_key, _) in world.imports.iter() {
128 if let WorldKey::Interface(iface_id) = import_key {
129 let iface = &wit_bindgen_cfg.resolve.interfaces[*iface_id];
131
132 if iface
138 .package
139 .map(|p| &wit_bindgen_cfg.resolve.packages[p].name)
140 .is_some_and(is_ignored_invocation_handler_pkg)
141 {
142 continue;
143 }
144
145 for (iface_fn_name, iface_fn) in iface.functions.iter() {
161 debug!("processing imported interface function: [{iface_fn_name}]");
162 imported_iface_invocation_methods.push(
163 translate_import_fn_for_lattice(iface, iface_fn_name, iface_fn, &cfg)
164 .expect("failed to translate export fn"),
165 );
166 }
167 }
168 }
169 }
170
171 let bindgen_tokens: TokenStream =
174 expand_wasmtime_component(wit_bindgen_cfg).unwrap_or_else(syn::Error::into_compile_error);
175
176 let mut bindgen_ast: syn::File =
179 syn::parse2(bindgen_tokens).expect("failed to parse wit-bindgen generated code as file");
180
181 let mut visitor = WitBindgenOutputVisitor::new(&cfg);
184 visitor.visit_file_mut(&mut bindgen_ast);
185
186 let methods_by_iface = build_lattice_methods_by_wit_interface(
189 &visitor.serde_extended_structs,
190 &visitor.type_lookup,
191 &visitor.export_trait_methods,
192 &cfg,
193 )
194 .expect("failed to build lattice methods from WIT interfaces");
195
196 let impl_struct_name = Ident::new_raw(cfg.impl_struct.as_str(), Span::call_site());
198
199 let mut interface_dispatch_wrpc_match_arms: Vec<TokenStream> = Vec::new();
201 let mut iface_tokens = TokenStream::new();
202
203 for (wit_iface_name, methods) in methods_by_iface.iter() {
210 let wit_iface = Ident::new(wit_iface_name, Span::call_site());
212
213 let operation_names = methods
216 .clone()
217 .into_iter()
218 .map(|lm| lm.operation_name)
219 .collect::<Vec<LitStr>>();
220
221 let func_names = methods
223 .clone()
224 .into_iter()
225 .map(|lm| lm.func_name)
226 .collect::<Vec<Ident>>();
227
228 let invocation_args_with_types = methods
233 .clone()
234 .into_iter()
235 .map(|lm| {
236 let arg_tokens = lm
237 .invocation_args
238 .iter()
239 .map(|(ident, ty)| quote!(#ident: #ty))
240 .collect::<Vec<TokenStream>>();
241 quote::quote!(#( #arg_tokens ),*)
242 })
243 .collect::<Vec<TokenStream>>();
244
245 let invocation_returns = methods
247 .clone()
248 .into_iter()
249 .map(|lm| lm.invocation_return)
250 .collect::<Vec<ReturnType>>();
251
252 iface_tokens.append_all(quote!(
258 #[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
259 pub trait #wit_iface {
260 fn contract_id() -> &'static str {
261 #contract_ident
262 }
263
264 #(
265 async fn #func_names (
266 &self,
267 ctx: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Context,
268 #invocation_args_with_types
269 ) #invocation_returns;
270 )*
271 }
272 ));
274
275 let (wrpc_input_parsing_statements, post_self_args, result_encode_tokens) = methods.clone().into_iter().fold(
283 (Vec::<TokenStream>::new(), Vec::<TokenStream>::new(), Vec::<TokenStream>::new()),
284 |mut acc, lm| {
285 let mut input_decoding_lines = Vec::<TokenStream>::new();
298
299 for (arg_name, arg_type) in lm.invocation_args.iter() {
303 let arg_name_lit = LitStr::new(&arg_name.to_string(), Span::call_site());
304 let arg_ty = arg_type.to_token_stream();
305 input_decoding_lines.push(quote::quote!(
306 let mut #arg_name = ::wasmcloud_provider_wit_bindgen::deps::bytes::BytesMut::new();
307 params
308 .pop()
309 .ok_or_else(|| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("missing expected parameter [{}]", #arg_name_lit)))?
310 .encode(&mut #arg_name)
311 .await
312 .map_err(|e| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("failed to encode parameter [{}]: {e}", #arg_name_lit)))?;
313 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)
314 .await
315 .map_err(|e| ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(format!("failed to receive parameter [{}]: {e}", #arg_name_lit)))?;
316 ));
317 }
318 acc.0.push(quote::quote!(#( #input_decoding_lines );*));
319
320 let arg_idents = vec![Ident::new("ctx", Span::call_site())]
323 .into_iter()
324 .chain(lm.invocation_args.iter().map(|(name, _)| name.clone()))
325 .collect::<Vec<Ident>>();
326 acc.1.push(quote!(#( #arg_idents ),*));
327
328 acc.2.push(match lm.invocation_return {
331 syn::ReturnType::Type(_, _) => {
332 quote!(result
333 .encode(&mut res)
334 .await
335 .map_err(|e| {
336 ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(
337 format!("failed to encode result of operation [{operation}]: {e}")
338 )
339 })?)
340 }
341
342 syn::ReturnType::Default => {
344 quote!(result
345 .encode(&mut res)
346 .await
347 .map_err(|e| {
348 ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Unexpected(
349 format!("failed to encode result of operation [{operation}]: {e}")
350 )
351 })?)
352 },
353 });
354
355 acc
356 },
357 );
358
359 interface_dispatch_wrpc_match_arms.push(quote!(
360 #(
361 operation @ #operation_names => {
362 #wrpc_input_parsing_statements
363 let result = #wit_iface::#func_names(
364 self,
365 #post_self_args
366 )
367 .await;
368
369 let mut res = ::wasmcloud_provider_wit_bindgen::deps::bytes::BytesMut::new();
370 #result_encode_tokens;
371 Ok(res.to_vec())
372 }
373 )*
374 ));
375 }
376
377 let types: Vec<TokenStream> = visitor
379 .type_lookup
380 .iter()
381 .filter_map(|(_, (_, ty))| {
382 if visitor
385 .serde_extended_structs
386 .contains_key(&ty.ident.to_string())
387 || visitor
388 .serde_extended_enums
389 .contains_key(&ty.ident.to_string())
390 {
391 None
392 } else {
393 Some(ty.to_token_stream())
394 }
395 })
396 .collect();
397
398 let structs: Vec<TokenStream> = visitor
400 .serde_extended_structs
401 .iter()
402 .map(|(_, (_, s))| s.to_token_stream())
403 .collect();
404
405 let enums: Vec<TokenStream> = visitor
407 .serde_extended_enums
408 .iter()
409 .map(|(_, (_, s))| s.to_token_stream())
410 .collect();
411
412 let wrpc_impl_tokens = build_wrpc_impls(&impl_struct_name, &wit_bindgen_cfg.resolve)
414 .expect("failed to build provider-sdk wrpc implementation");
415
416 let tokens = quote!(
418 #iface_tokens
420 #(
424 #types
425 )*
426 #(
430 #structs
431 )*
432 #(
436 #enums
437 )*
438 #[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
447 trait WasmcloudCapabilityProvider {
448 async fn receive_link_config_as_source(
449 &self,
450 link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
451 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
452 Ok(())
453 }
454
455 async fn receive_link_config_as_target(
456 &self,
457 link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
458 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
459 Ok(())
460 }
461
462 async fn delete_link(
463 &self,
464 actor_id: &str
465 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
466 Ok(())
467 }
468
469 async fn shutdown(&self) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
470 Ok(())
471 }
472 }
473
474 #[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
479 impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderHandler for #impl_struct_name {
480 async fn receive_link_config_as_source(
481 &self,
482 link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
483 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
484 WasmcloudCapabilityProvider::receive_link_config_as_source(self, link_config).await
485 }
486
487 async fn receive_link_config_as_target(
488 &self,
489 link_config: impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::LinkConfig
490 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
491 WasmcloudCapabilityProvider::receive_link_config_as_target(self, link_config).await
492
493 }
494
495 async fn delete_link(
496 &self,
497 actor_id: &str
498 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
499 WasmcloudCapabilityProvider::delete_link(self, actor_id).await
500 }
501
502 async fn shutdown(&self) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::ProviderOperationResult<()> {
503 WasmcloudCapabilityProvider::shutdown(self).await
504 }
505 }
506
507 impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Provider for #impl_struct_name {}
510
511 pub struct InvocationHandler {
516 wrpc_client: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::core::wrpc::Client
517 }
518
519 impl InvocationHandler {
520 pub fn new(target: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::core::ComponentId) -> Self {
521 let connection = ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::get_connection();
522 Self { wrpc_client: connection.get_wrpc_client(&target) }
527 }
528
529 #(
530 #imported_iface_invocation_methods
531 )*
532 }
533
534 #wrpc_impl_tokens
535
536 #[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
537 impl ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::WrpcDispatch for #impl_struct_name {
538 async fn dispatch_wrpc_dynamic<'a>(
539 &'a self,
540 ctx: ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::Context,
541 operation: String,
542 mut params: Vec<::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::Value>,
543 ) -> ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationResult<Vec<u8>> {
544 use ::wasmcloud_provider_wit_bindgen::deps::wrpc_transport::{Encode, Receive};
545 use ::wasmcloud_provider_wit_bindgen::deps::anyhow::Context as _;
546 match operation.as_str() {
547 #(
548 #interface_dispatch_wrpc_match_arms
549 )*
550 _ => Err(::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::InvocationError::Malformed(format!(
551 "Invalid operation name [{operation}]"
552 )).into())
553 }
554 }
555 }
556 );
557
558 tokens.into()
559}
560
561fn build_lattice_methods_by_wit_interface(
564 struct_lookup: &StructLookup,
565 type_lookup: &TypeLookup,
566 export_trait_methods: &HashMap<WitInterfacePath, Vec<ImplItemFn>>,
567 bindgen_cfg: &ProviderBindgenConfig,
568) -> anyhow::Result<HashMap<WitTraitName, Vec<ExportedLatticeMethod>>> {
569 let mut methods_by_name: HashMap<WitInterfacePath, Vec<ExportedLatticeMethod>> = HashMap::new();
570
571 for (wit_iface_name, funcs) in export_trait_methods.iter() {
574 for trait_method in funcs.iter() {
575 let wit_operation = match wit_iface_name.split('.').collect::<Vec<&str>>()[..] {
577 [wit_ns, wit_pkg, iface] => {
578 format!(
579 "{}:{}/{}.{}",
580 wit_ns.to_kebab_case(),
581 wit_pkg.to_kebab_case(),
582 iface.to_kebab_case(),
583 trait_method.sig.ident.to_string().to_kebab_case()
584 )
585 }
586 _ => bail!("unexpected interface path, expected 3 components"),
587 };
588 let operation_name = LitStr::new(&wit_operation, trait_method.sig.ident.span());
589
590 let lattice_method = translate_export_fn_for_lattice(
592 bindgen_cfg,
593 operation_name,
594 trait_method,
595 struct_lookup,
596 type_lookup,
597 )?;
598
599 let wit_iface_upper_camel = wit_iface_name
601 .split('.')
602 .map(|v| v.to_upper_camel_case())
603 .collect::<String>();
604
605 methods_by_name
609 .entry(wit_iface_upper_camel)
610 .or_default()
611 .push(lattice_method);
612 }
613 }
614 Ok(methods_by_name)
615}
616
617fn is_ignored_invocation_handler_pkg(pkg: &wit_parser::PackageName) -> bool {
619 matches!(
620 (pkg.namespace.as_ref(), pkg.name.as_ref()),
621 ("wasmcloud", "bus") | ("wasi", "io")
622 )
623}
624
625fn build_wrpc_impls(impl_struct_name: &Ident, resolve: &Resolve) -> anyhow::Result<TokenStream> {
627 let mapping = crate::wrpc::generate_wrpc_nats_subject_to_fn_mapping(resolve)
628 .context("failed to generate wrpc NATS subject mappings")?;
629
630 let mut insertion_lines: Vec<TokenStream> = Vec::new();
633 for crate::wrpc::WrpcExport {
634 wit_ns,
635 wit_pkg,
636 wit_iface,
637 wit_iface_fn,
638 types,
639 } in mapping.into_iter()
640 {
641 let wit_ns = LitStr::new(&wit_ns, Span::call_site());
642 let wit_pkg = LitStr::new(&wit_pkg, Span::call_site());
643 let wit_iface = LitStr::new(&wit_iface, Span::call_site());
644 let wit_iface_fn = LitStr::new(&wit_iface_fn, Span::call_site());
645 let world_key_name = LitStr::new(&types.0, Span::call_site());
646 let function_name = LitStr::new(&types.1, Span::call_site());
647 let dynamic_fn = LitStr::new(
648 &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}]")?,
649 Span::call_site(),
650 );
651
652 insertion_lines.push(quote!(
653 mapping.insert(
654 format!("{lattice_name}.{component_id}.wrpc.{wrpc_version}.{}:{}/{}.{}", #wit_ns, #wit_pkg, #wit_iface, #wit_iface_fn),
655 (#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")),
656 );
657 ));
658 }
659
660 let tokens = quote!(
662 use ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::{WrpcNats, WrpcNatsSubject, WorldKeyName, WitFunction};
663 use ::wasmcloud_provider_wit_bindgen::deps::wasmcloud_provider_sdk::error::ProviderInitResult;
664
665 #[::wasmcloud_provider_wit_bindgen::deps::async_trait::async_trait]
666 impl WrpcNats for #impl_struct_name {
667 async fn incoming_wrpc_invocations_by_subject(
668 &self,
669 lattice_name: impl AsRef<str> + Send,
670 component_id: impl AsRef<str> + Send,
671 wrpc_version: impl AsRef<str> + Send,
672 ) -> ProviderInitResult<
673 ::std::collections::HashMap<WrpcNatsSubject, (WorldKeyName, WitFunction, ::wasmcloud_provider_wit_bindgen::deps::wrpc_types::DynamicFunction)>
674 > {
675 let lattice_name = lattice_name.as_ref();
676 let wrpc_version = wrpc_version.as_ref();
677 let component_id = component_id.as_ref();
678 let mut mapping = ::std::collections::HashMap::new();
679 #(
680 #insertion_lines
681 )*
682 Ok(mapping)
683 }
684 }
685 );
686
687 Ok(tokens)
688}
689
690#[cfg(test)]
691mod tests {
692 use std::collections::HashMap;
693
694 use anyhow::Context;
695 use proc_macro2::{Span, TokenTree};
696 use quote::quote;
697 use syn::{parse_quote, ImplItemFn, LitStr, ReturnType};
698
699 use crate::wit::{extract_witified_map, translate_export_fn_for_lattice};
700 use crate::ProviderBindgenConfig;
701
702 #[test]
704 fn parse_witified_map_type() -> anyhow::Result<()> {
705 extract_witified_map(
706 "e!(Vec<(String, String)>)
707 .into_iter()
708 .collect::<Vec<TokenTree>>(),
709 )
710 .context("failed to parse WIT-ified map type Vec<(String, String)>")?;
711 Ok(())
712 }
713
714 #[test]
716 fn parse_witified_map_in_fn() -> anyhow::Result<()> {
717 let trait_fn: ImplItemFn = parse_quote!(
718 fn baz(test_map: Vec<(String, String)>) {}
719 );
720 let bindgen_cfg = ProviderBindgenConfig {
721 impl_struct: "None".into(),
722 contract: "wasmcloud:test".into(),
723 wit_ns: Some("test".into()),
724 wit_pkg: Some("foo".into()),
725 exposed_interface_allow_list: Default::default(),
726 exposed_interface_deny_list: Default::default(),
727 wit_bindgen_cfg: None, replace_witified_maps: true,
729 };
730 let operation_name = LitStr::new("wasmcloud:test/test.foo", Span::call_site());
731 let lm = translate_export_fn_for_lattice(
732 &bindgen_cfg,
733 operation_name.clone(),
734 &trait_fn,
735 &HashMap::new(), &HashMap::new(), )?;
738
739 assert_eq!(lm.operation_name, operation_name);
740 assert_eq!(lm.invocation_args.len(), 1);
741 assert_eq!(
742 lm.invocation_return,
743 syn::parse2::<ReturnType>(quote::quote!())?
744 );
745
746 Ok(())
747 }
748}