60·高级

自定义 derive 宏实战

自定义 derive 宏实战

学习目标

  1. 完整实现一个 derive 宏
  2. 理解 syn 的 AST 结构
  3. 掌握代码生成

核心概念

示例:Builder 宏

// 使用方
#[derive(Builder)]
struct Server {
    host: String,
    port: u16,
    max_connections: usize,
}

// 生成:
// struct ServerBuilder {
//     host: Option<String>,
//     port: Option<u16>,
//     max_connections: Option<usize>,
// }
// impl ServerBuilder { ... }

实现

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, DeriveInput, Data, Fields};

#[proc_macro_derive(Builder)]
pub fn derive_builder(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let name = &input.ident;
    let builder_name = format!("{}Builder", name);

    let fields = match &input.data {
        Data::Struct(data) => match &data.fields {
            Fields::Named(fields) => &fields.named,
            _ => panic!("Only named fields supported"),
        },
        _ => panic!("Only structs supported"),
    };

    let builder_fields = fields.iter().map(|f| {
        let name = &f.ident;
        let ty = &f.ty;
        quote! { #name: Option<#ty> }
    });

    let builder_methods = fields.iter().map(|f| {
        let name = &f.ident;
        let ty = &f.ty;
        quote! {
            pub fn #name(mut self, value: #ty) -> Self {
                self.#name = Some(value);
                self
            }
        }
    });

    let build_fields = fields.iter().map(|f| {
        let name = &f.ident;
        quote! {
            #name: self.#name.ok_or(format!("{} not set", stringify!(#name)))?
        }
    });

    let builder_ident = syn::Ident::new(&builder_name, name.span());

    let expanded = quote! {
        struct #builder_ident {
            #(#builder_fields,)*
        }

        impl #builder_ident {
            #(#builder_methods)*

            pub fn build(self) -> Result<#name, String> {
                Ok(#name {
                    #(#build_fields,)*
                })
            }
        }

        impl #name {
            fn builder() -> #builder_ident {
                #builder_ident {
                    #(#name: None,)*
                }
            }
        }
    };

    TokenStream::from(expanded)
}

使用

#[derive(Builder)]
struct Config {
    host: String,
    port: u16,
}

fn main() {
    let config = Config::builder()
        .host("localhost".to_string())
        .port(8080)
        .build()
        .unwrap();

    println!("{}:{}", config.host, config.port);
}

小结

工具用途
syn::parse_macro_input!解析输入 TokenStream
quote! { }生成代码
#variable在 quote 中插入变量
#(#items)*重复生成

练习编辑器

rust
Loading...