Blocktree/crates/btproto/src/generation.rs

use std::collections::HashMap;

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

use crate::{
    case_convert::CaseConvert,
    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 variants: Vec<_> = msg_lookup.msg_iter().map(|msg| msg.msg_name()).collect();
        let variant_names_map: HashMap<&Ident, String> = variants
            .iter()
            .map(|variant| (variant.as_ref(), variant.to_string()))
            .collect();
        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 enum_kinds_name = format_ident!("{}MsgKinds", self.def().name_def.name);
        let name_decl_vec: Vec<(Ident, TokenStream)> = variants
            .iter()
            .map(|variant| {
                let variant_str = variant_names_map.get(variant.as_ref()).unwrap();
                let name_ident =
                    format_ident!("{}_NAME", variant_str.pascal_to_snake().to_uppercase());
                let decl = quote! {
                    static #name_ident: Lazy<Arc<String>> = Lazy::new(|| {
                        Arc::new(#variant_str.into())
                    });
                };
                (name_ident, decl)
            })
            .collect();
        let name_decls = name_decl_vec.iter().map(|(_, decl)| decl);
        let name_idents = name_decl_vec.iter().map(|(ident, _)| ident);
        let send_impl = if all_replies {
            quote! {}
        } else {
            quote! {
                impl ::btrun::model::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 {
                #( #variants(#msg_types) ),*
            }


            #[derive(Clone)]
            #[allow(dead_code)]
            pub enum #enum_kinds_name {
                #( #variants ),*
            }

            impl ::btrun::model::Named for #enum_kinds_name {
                fn name(&self) -> ::std::sync::Arc<String> {
                    use ::btrun::model::Lazy;
                    use ::std::sync::Arc;
                    #( #name_decls )*
                    match self {
                        #( Self::#variants => #name_idents.clone() ),*
                    }
                }
            }

            impl Copy for #enum_kinds_name {}

            impl ::btrun::model::Named for #enum_name {
                fn name(&self) -> ::std::sync::Arc<String> {
                    match self {
                        #( Self::#variants(_) => #enum_kinds_name::#variants.name()),*
                    }
                }
            }

            impl ::btrun::model::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(),
                self.actor_lookup()
                    .actor_with_init_state(state.name())
                    .is_some(),
            )
        });
        let mut tokens = TokenStream::new();
        for (trait_ident, methods, is_init_state) in traits {
            let method_tokens = methods.map(|x| x.generate_tokens());
            let actor_impl_method = if is_init_state {
                quote! {
                    #[doc = "The name of the implementation for the actor this state is a part of."]
                    fn actor_impl() -> ::std::sync::Arc<String>;
                }
            } else {
                quote! {}
            };
            let trait_str = format!("{trait_ident}");
            quote! {
                pub trait #trait_ident : Send + Sync + Sized {
                    #actor_impl_method
                    #( #method_tokens )*

                    fn state_name() -> ::std::sync::Arc<String> {
                        use ::btrun::model::Lazy;
                        use ::std::sync::Arc;
                        static STATE_NAME: Lazy<Arc<String>> = Lazy::new(|| {
                            Arc::new(#trait_str.into())
                        });
                        STATE_NAME.clone()
                    }
                }
            }
            .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 = ::btrun::model::TransResult<Self, ( #( #output_types ),* )>
            >;
            fn #method_ident(self #( , #msg_args )*) -> Self::#future_name;
        }
    }
}