【社区投稿】给Rust的Struct自动实现trait

给Rust的Struct自动实现trait

我们通常使用

go 复制代码

#[derive(Clone, Debug)]

这样的方式给struct自动实现相应的trait，从而让struct具备某些特性，但是如果我们想让编译器给struct自动实现自己定义的trait要怎么办？

首先我们需要有一个trait，假设如下面的定义：

go 复制代码

pub trait Printable {
    pub fn print_me(&self);
}

我们定义这个trait给struct赋予一个行为是逐行打印struct的所有Field。当然如果是自己实现肯定是可以凭空乱写的，那么我们可以和Debug一样，在 derive 中让编译器自动添加默认的实现。

首先需要给crate添加一个子crate：

go 复制代码

cargo new --lib printable

然后在当前crate的 Cargo.toml 中添加依赖

go 复制代码

[workspace]
members = [
    ".",
    "printable"
]
[dependencies]
printable = { version = "*", path = "printable"}

在printable 的 Cargo.toml 里还需要添加依赖

go 复制代码

[dependencies]
syn = { version = "1.0", features = ["full"] }
quote = "1.0"
proc-macro2 = "1.0.51"

我们需要这三个crate来简化代码生成的工作，这里proc-macro2提供了自动实现宏的功能，syn用来解析结构体，quote用来输出TokenStream。

在 printable 的lib.rs 文件中

go 复制代码

#[proc_macro_derive(Printable)]
pub fn print_info_derive(input: TokenStream) -> TokenStream {
}

我们在函数 print_info_derive 中输出的TokenStream，就会在编译时动态注入到struct中，这里参数input就是struct本身的代码流。

我们通过解析input就可以分析出 Struct的名字，Field列表，所有Field的名字，类型.....

下面是简化后的代码：

go 复制代码

#[proc_macro_derive(Printable)]
pub fn print_info_derive(input: TokenStream) -> TokenStream {
    let struct_name = to_snake_case(input.ident.to_string().as_str());
        
    let fields = match input.data.clone() {
        syn::Data::Struct(data) => data.fields,
        _ => panic!("Only structs are supported"),
    };
    let fields_name: Vec<Ident> = fields.iter().map(|field| {
        field.ident.as_ref().unwrap().clone()
    }).collect();
}

之后我们就需要构建输出的代码流，这里使用 quote! 这个宏来实现。

go 复制代码

#[proc_macro_derive(Printable)]
pub fn print_info_derive(input: TokenStream) -> TokenStream {
    let struct_name = to_snake_case(input.ident.to_string().as_str());
        
    let fields = match input.data.clone() {
        syn::Data::Struct(data) => data.fields,
        _ => panic!("Only structs are supported"),
    };
    let fields_name: Vec<Ident> = fields.iter().map(|field| {
        field.ident.as_ref().unwrap().clone()
    }).collect();
    let output_token = quote! {
        impl Printable for #struct_name {
            pub fn print_me(&self) {
                //这里添加逐行打印Field的代码，因为quote里本来就是在输出代码流
                //所以不能直接访问fields_name，比如循环之类的，所以我们这里需要
                //把生成这部分代码提取到函数外
            }
        }
    }
    output_token.into()
}

为了简单演示我们就使用一个函数来实现：

go 复制代码

fn gen_print(fileds: Vec<Ident>) -> TokenStream2 {
    let print_stmts =fields.iter().map(|field| {
        quote! {
            println!("field:{}", &self.#field);
        }
    });
    quote!{
        #(#print_stmts)*
    }
}

最后组装一下，lib.rs 的代码如下：

go 复制代码

#[proc_macro_derive(Printable)]
pub fn print_info_derive(input: TokenStream) -> TokenStream {
    let struct_name = to_snake_case(input.ident.to_string().as_str());
        
    let fields = match input.data.clone() {
        syn::Data::Struct(data) => data.fields,
        _ => panic!("Only structs are supported"),
    };
    let fields_name: Vec<Ident> = fields.iter().map(|field| {
        field.ident.as_ref().unwrap().clone()
    }).collect();
    let print_code = gen_print(fields_name);
    let output_token = quote! {
        impl Printable for #struct_name {
            pub fn print_me(&self) {
                #output_token
            }
        }
    }
    output_token.into()
}

如此这般一通操作后，我们随便一定一个Struct：

go 复制代码

#[derive(Debug, Printable)]
struct TestTb {
    id: String,
    name: String,
    ts: i32
}

就可以

go 复制代码

TestTb{id: "123".to_string(), name: "alex", ts: 111}.print_me();

就可以逐行打印出所有的Field了。

那么灵活的使用这个玩法，我们可以根据Struct的Field，自动生成 insert, update, delete的SQL也是可以的。给每个Field自动生成getter，setter方法...... （这个Java味太浓了，だめ）

研究这个是为了给 sqlx 增加一个自动生成insert，update，delete方法的增强，因为不喜欢写超长的insert和update语句。

Amusez-vous tous！