Blocktree/crates/btproto/src/generation.rs

use proc_macro2::TokenStream;
use quote::{format_ident, quote, ToTokens};

use crate::model::{MethodModel, ProtocolModel};

impl ToTokens for ProtocolModel {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        tokens.extend(self.generate_message_enum());
        tokens.extend(self.generate_state_traits());
    }
}

impl ProtocolModel {
    fn generate_message_enum(&self) -> TokenStream {
        let msg_lookup = self.msg_lookup();
        let get_variants = || msg_lookup.msg_iter().map(|msg| msg.msg_name());
        let variants0 = get_variants();
        let variants1 = get_variants();
        let variant_names = get_variants().map(|variant| variant.to_string());
        let msg_types = msg_lookup.msg_iter().map(|msg| msg.msg_type());
        let all_replies = msg_lookup.msg_iter().all(|msg| msg.is_reply());
        let enum_name = format_ident!("{}Msgs", self.def().name_def.name);
        let send_impl = if all_replies {
            quote! {}
        } else {
            quote! {
                impl ::btrun::SendMsg for #enum_name {}
            }
        };
        let proto_name = &self.def().name_def.name;
        let doc_comment = format!("Message type for the {proto_name} protocol.");
        quote! {
            #[doc = #doc_comment]
            #[derive(::serde::Serialize, ::serde::Deserialize)]
            pub enum #enum_name {
                #( #variants0(#msg_types) ),*
            }

            impl #enum_name {
                pub fn name(&self) -> &'static str {
                    match self {
                        #( Self::#variants1(_) => #variant_names),*
                    }
                }
            }

            impl ::btrun::CallMsg for #enum_name {
                type Reply = Self;
            }

            #send_impl
        }
    }

    fn generate_state_traits(&self) -> TokenStream {
        let traits = self
            .states_iter()
            .map(|state| (state.name(), state.methods().values()));
        let mut tokens = TokenStream::new();
        for (trait_ident, methods) in traits {
            let method_tokens = methods.map(|x| x.generate_tokens());
            quote! {
                pub trait #trait_ident : Send + Sync {
                    #( #method_tokens )*
                }
            }
            .to_tokens(&mut tokens);
        }
        tokens
    }
}

impl MethodModel {
    /// Generates the tokens for the code which implements this transition.
    fn generate_tokens(&self) -> TokenStream {
        let method_ident = self.name().as_ref();
        let msg_args = self.inputs().iter();
        let output_decls = self.outputs().iter().flat_map(|output| output.decl());
        let output_types = self.outputs().iter().flat_map(|output| output.type_name());
        let future_name = self.future();
        quote! {
            #( #output_decls )*
            type #future_name: Send + ::std::future::Future<Output = Result<( #( #output_types ),* )>>;
            fn #method_ident(self #( , #msg_args )*) -> Self::#future_name;
        }
    }
}