hsnet_rpc_macro/

lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, Attribute, FnArg, ItemTrait, Meta, Pat, TraitItem};

/// RPC trait 宏
///
/// 用法：
/// ```ignore
/// #[rpc(server)]
/// pub trait ControlRpc {
///     #[method(name = "wg_start")]
///     async fn wg_start(&self, config: &WgConfig) -> Result<String, RpcError>;
/// }
/// ```
#[proc_macro_attribute]
pub fn rpc(_args: TokenStream, input: TokenStream) -> TokenStream {
    let input_trait = parse_macro_input!(input as ItemTrait);

    let trait_name = &input_trait.ident;
    let trait_vis = &input_trait.vis;

    // 提取所有 RPC 方法
    let mut rpc_methods = Vec::new();

    for item in &input_trait.items {
        if let TraitItem::Fn(method) = item {
            // 查找 #[method] 属性
            let method_attr = method.attrs.iter().find(|attr| attr.path().is_ident("method"));

            if let Some(attr) = method_attr {
                // 解析 #[method(name = "xxx")]
                let rpc_name = parse_method_name(attr);
                if rpc_name.is_empty() {
                    continue;
                }

                let method_name = &method.sig.ident;
                let method_inputs = &method.sig.inputs;

                // 提取参数（跳过 &self）
                let params: Vec<_> = method_inputs
                    .iter()
                    .skip(1) // 跳过 &self
                    .filter_map(|arg| {
                        if let FnArg::Typed(pat_type) = arg {
                            if let Pat::Ident(pat_ident) = &*pat_type.pat {
                                let param_name = &pat_ident.ident;
                                let param_type = &pat_type.ty;
                                return Some((param_name.clone(), param_type.clone()));
                            }
                        }
                        None
                    })
                    .collect();

                // 提取 cfg 属性
                let cfg_attrs: Vec<_> = method
                    .attrs
                    .iter()
                    .filter(|attr| attr.path().is_ident("cfg"))
                    .cloned()
                    .collect();

                rpc_methods.push(RpcMethod {
                    rpc_name,
                    method_name: method_name.clone(),
                    params,
                    cfg_attrs,
                });
            }
        }
    }

    // 生成 handler 注册代码
    let handlers = rpc_methods.iter().map(|method| {
        let rpc_name = &method.rpc_name;
        let method_name = &method.method_name;
        let cfg_attrs = &method.cfg_attrs;

        if method.params.is_empty() {
            // 无参数方法
            quote! {
                #(#cfg_attrs)*
                {
                    server = server.handle(#rpc_name, {
                        let service = service.clone();
                        move |_: ()| {
                            let service = service.clone();
                            async move { service.#method_name().await }
                        }
                    });
                }
            }
        } else if method.params.len() == 1 {
            // 单参数方法 - 直接传递值，不加引用
            let param_name = &method.params[0].0;
            let param_type = &method.params[0].1;

            quote! {
                #(#cfg_attrs)*
                {
                    server = server.handle(#rpc_name, {
                        let service = service.clone();
                        move |#param_name: #param_type| {
                            let service = service.clone();
                            async move { service.#method_name(#param_name).await }
                        }
                    });
                }
            }
        } else {
            // 多参数 - 暂不支持
            quote! {
                compile_error!("Multiple parameters not yet supported");
            }
        }
    });

    // 过滤掉 #[rpc] 属性，保留其他属性
    let trait_attrs: Vec<_> = input_trait
        .attrs
        .iter()
        .filter(|attr| !attr.path().is_ident("rpc"))
        .collect();

    // 清理 trait items，去掉所有 #[method] 属性
    let cleaned_items: Vec<_> = input_trait
        .items
        .iter()
        .map(|item| {
            if let TraitItem::Fn(method) = item {
                let mut method = method.clone();
                method.attrs.retain(|attr| !attr.path().is_ident("method"));
                TraitItem::Fn(method)
            } else {
                item.clone()
            }
        })
        .collect();

    let trait_generics = &input_trait.generics;

    // 生成最终代码
    let expanded = quote! {
        // 保留原始 trait 定义（去掉 #[rpc] 和 #[method] 属性）
        #(#trait_attrs)*
        #trait_vis trait #trait_name #trait_generics {
            #(#cleaned_items)*
        }

        // 生成扩展 trait
        #trait_vis trait IntoRpcServer {
            fn into_rpc(self: std::sync::Arc<Self>) -> hsnet_rpc::RpcServer;
        }

        // 为所有实现了原 trait 的类型实现 IntoRpcServer
        impl<T: #trait_name + Send + Sync + 'static> IntoRpcServer for T {
            fn into_rpc(self: std::sync::Arc<Self>) -> hsnet_rpc::RpcServer {
                let service = self;
                let mut server = hsnet_rpc::RpcServer::new();

                #(#handlers)*

                server
            }
        }
    };

    TokenStream::from(expanded)
}

// 辅助函数：解析 #[method(name = "xxx")]
fn parse_method_name(attr: &Attribute) -> String {
    if let Meta::List(meta_list) = &attr.meta {
        let tokens_str = meta_list.tokens.to_string();

        // 简单解析 name = "xxx"
        if let Some(start) = tokens_str.find('"') {
            if let Some(end) = tokens_str[start + 1..].find('"') {
                return tokens_str[start + 1..start + 1 + end].to_string();
            }
        }
    }
    String::new()
}

struct RpcMethod {
    rpc_name: String,
    method_name: syn::Ident,
    params: Vec<(syn::Ident, Box<syn::Type>)>,
    cfg_attrs: Vec<Attribute>,
}