Rust 宏魔法系列 - 派生宏

简介

宏是一种生成代码的手段，常用于 DSL 设计, 比如 jsx 本质上就是一种 DSL, 通过 babel 等工具转换成 js 再运行, 目的是为了降低编码的复杂度

派生宏是 Rust 宏的一种, 用来于生成额外的代码, 例如 #[derive(Clone)] 里面的 Clone 就是一种派生宏, 会自动生成 impl Clone for XXX 这样的模板代码

本文通过一个例子来介绍如何编写派生宏

写一个派生宏

手写一个派生宏 Optional

将结构体的所有字段转换成可选的字段，生成一个 Optional<Struct>

例如原本的结构体如下所示

struct Picea {
    id: usize,
    subs: Option<Vec<Picea>>,
    descriptions: String,
}

期望的派生宏形式是这样的

#[derive(Optional)]
struct Picea {
    id: usize,
    subs: Option<Vec<Picea>>,
    descriptions: String,
}

加上派生宏之后,生成如下的代码

struct OptionalPicea {
    id: Option<usize>,
    subs: Option<Vec<Picea>>,
    descriptions: Option<String>,
}

还可以加上一点属性,例如 skip, rename

#[derive(Optional)]
struct Picea {
    #[optional(rename = value)]  // 属性名改成 value
    id: usize,
    subs: Option<Vec<Picea>>,
    #[optional(skip)] // 有了 skip, description  就不会被 Option
    descriptions: String,
}

会生成这样的代码

struct OptionalPicea {
    value: Option<usize>,
    subs: Option<Vec<Picea>>,
    descriptions: String,
}

开始编码

准备工作

先创建一个lib, 在 cargo.toml 里面加入

[lib]
proc-macro = true

[dependencies]
quote = "1.0.35"
syn = "2.0.52"

lib.rs 的内容替换成如下所示

use proc_macro::TokenStream;

#[proc_macro_derive(Optional, attributes(optional))]
pub fn optional(input: TokenStream) -> TokenStream {
    input.clone()
}

可以用自己编写 test 文件进行最终的测试, 或者在另外一个项目里面引入该 crate 进行测试, 下面是测试的代码

// 测试的代码
#[derive(Optional)]
struct Picea {
    #[optional(rename = value)]
    id: usize,
    subs: Option<Vec<Picea>>,
    #[optional(skip)]
    descriptions: String,
}

实现部分

所谓的 TokenStream 就是 rust 的代码,可以理解成一个数组, 每个元素都是一个树, 里面包含TokenTree, 不过我们不会直接操作 TokenStream, 一般都通过 syn 和 quote 来进行解析和展开的工作

syn 和 quote 的详细用法可以直接去看他们的文档

首先通过 sync 对 input 进行解析, 得到一个 DeriveInput

use syn::{parse_macro_input, DeriveInput};

    let input = parse_macro_input!(input as DeriveInput);

其中 DeriveInput 包含了 struct 的所有信息, 包括 ident(结构体名字) 和里面的字段 , 派生宏只能被用于 结构体,枚举,Union

use syn::{parse_macro_input, Data, DeriveInput};

    let ident = input.ident; // 结构体的名字, 在这里就是 "Picea"
    match input.data {
        Data::Struct(data) => {} // 结构体
        Data::Enum(data) => {} // 枚举
        Data::Union(data) => {} // Union
    };

因为这里我们只处理 Struct ,不处理其他类型, 我们可以只关心 Data::Struct, 其他情况直接报错

拼接我们要生成的结构体名, 这里生成一个新的 ident (OptionalPicea)

    let optional_struct_name = &format!("Optional{}", ident); // OptionalPicea
    let optional_struct_ident = Ident::new(optional_struct_name, ident.span());

使用 quote 进行代码展开,转换成 TokenStream

 quote::quote! {
    struct #optional_struct_ident {

    }
 }

上面这段代码会将 optional_struct_ident 替换成 OptionalPicea, 现在看一下整体的代码

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

#[proc_macro_derive(Optional, attributes(optional))]
pub fn optional(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let ident = input.ident;

    let Data::Struct(data_struct) = input.data else {
        // 不接受除了结构体之外的类型
        return syn::Error::new(ident.span(), "optional can only be applied to structs")
            .into_compile_error()
            .into();
    };

    let optional_struct_name = &format!("Optional{}", ident); // OptionalPicea
    let optional_struct_ident = Ident::new(optional_struct_name, ident.span());

    quote::quote! {
        struct #optional_struct_ident {

        }
    }
    .into()
}

上面的代码,会为 Picea 生成一个名为OptionalPicea的结构体, 我们下面迭代该结构体的每一个属性生成新的属性

    let fields: Vec<_> = data_struct
        .fields
        .iter()
        .map(|field| {
            // 对每个字段进行映射
            let ident = &field.ident;  // 属性名
            let ty = &field.ty; // 属性的类型
            quote::quote!(
                #ident: Option<#ty>
            )
        })
        .collect();

然后将 fields 放入结构体中

    quote::quote! {
        struct #optional_struct_ident {
            #(#fields,)* 
        }
    }

可以通过 cargo-expand 检查一下生成的代码, 如下所示

struct OptionalPicea {
    id: Option<usize>,
    subs: Option<Option<Vec<Picea>>>,
    descriptions: Option<String>,
}

可以看到中间的 sub 变成了 Option<Option<Vec<Picea>>>, 我们希望已经 Option 的字段不会再 Option, 就要多做一点判断

// 判断每一个属性的类型开头是不是 Option , 如果是  Option 开头就不处理,直接返回
if let Type::Path(path) = ty {  // Option<xxxx> 是  Path 形式的 Type
    let path = &path.path;
    let is_option = path
        .segments
        .last()
        .map(|segment| segment.ident == "Option")
        .unwrap_or(false);

    if is_option {
        return quote::quote!(
            #ident: #ty
        );
    }
}

展开结果如下

struct OptionalPicea {
    id: Option<usize>,
    subs: Option<Vec<Picea>>,
    descriptions: Option<String>,
}

将代码整理一下, 整体的代码如下所示

use proc_macro::TokenStream;
use syn::{parse_macro_input, Data, DeriveInput, Ident, Type};

#[proc_macro_derive(Optional, attributes(optional))]
pub fn optional(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let ident = input.ident;

    let Data::Struct(data_struct) = input.data else {
        // 不接受除了结构体之外的类型
        return syn::Error::new(ident.span(), "optional can only be applied to structs")
            .into_compile_error()
            .into();
    };

    let optional_struct_name = &format!("Optional{}", ident); // OptionalPicea
    let optional_struct_ident = Ident::new(optional_struct_name, ident.span());

    let fields: Vec<_> = data_struct
        .fields
        .iter()
        .map(|field| {
            // 对每个字段进行映射
            let ident = &field.ident;
            let ty = &field.ty;

            if is_ty_option(ty) {
                return quote::quote!(#ident: #ty);
            }

            quote::quote!(#ident: Option<#ty>)
        })
        .collect();

    quote::quote! {
        struct #optional_struct_ident {
            #(#fields,)*
        }
    }
    .into()
}

fn is_ty_option(ty: &Type) -> bool {
    let Type::Path(path) = ty else {
        return false;
    };
    let path = &path.path;
    path.segments
        .last()
        .map(|segment| segment.ident == "Option")
        .unwrap_or(false)
}

然后再看看 skip 字段, 我们希望标记了 #[optional(skip)] 的字段不进行 Option

在 attr 里面找到 optional 的字段

field.attrs.iter().find(|attr| attr.path().is_ident("optional"));

当为 optional(skip) 的时候, 才能跳过

attr.parse_args::<syn::Ident>()
                        .map(|ident| ident == "skip")
                        .unwrap_or(false)

这一块合并起来的逻辑为

let is_skip = field
    .attrs
    .iter()
    .find(|attr| attr.path().is_ident("optional"))
    .map(|attr| {
        attr.parse_args::<syn::Ident>()
            .map(|ident| ident == "skip")
            .unwrap_or(false)
    })
    .unwrap_or(false);

再看看展开的代码

struct OptionalPicea {
    id: Option<usize>,
    subs: Option<Vec<Picea>>,
    descriptions: String,
}

貌似是正确的, 但是把上面的代码改一改, 加入 rename 选项

#[derive(Optional)]
struct Picea {
    id: usize,
    subs: Option<Vec<Picea>>,
    #[optional(skip,rename = desc)]  // 将 rename 拼接在后面
    descriptions: String,
}

可以看到, skip 属性失效了, 这个时候, 这里我建议不要用 parse_args 进行解析了, 因为这个地方不做特殊的处理话,会连带 skip,rename = desc 一起进行解析

使用 parse_nested_meta 进行解析, 如下面的代码所示

let mut is_skip = false;

let attr = field
    .attrs
    .iter()
    .find(|attr| attr.path().is_ident("optional"));

if let Some(attr) = attr {
    let _ = attr.parse_nested_meta(|meta| {
        if meta.path.is_ident("skip") {
            is_skip = true; // 标记跳过
        } else if meta.path.is_ident("rename") {
            let renamed_ident = meta.value()?.parse::<syn::Ident>()?;
            ident = Some(renamed_ident);  // 更新  rename 的 属性
        }
        Ok(())
    });
}

得到结果如下

struct OptionalPicea {
    id: Option<usize>,
    subs: Option<Vec<Picea>>,
    desc: String,
}

最后看下整体的代码

use proc_macro::TokenStream;
use syn::{meta, parse_macro_input, Data, DeriveInput, Ident, Type};

#[proc_macro_derive(Optional, attributes(optional))]
pub fn optional(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let ident = input.ident;

    let Data::Struct(data_struct) = input.data else {
        // 不接受除了结构体之外的类型
        return syn::Error::new(ident.span(), "optional can only be applied to structs")
            .into_compile_error()
            .into();
    };

    let optional_struct_name = &format!("Optional{}", ident); // OptionalPicea
    let optional_struct_ident = Ident::new(optional_struct_name, ident.span());

    let fields: Vec<_> = data_struct
        .fields
        .iter()
        .map(|field| {
            // 对每个字段进行映射
            let mut ident = field.ident.clone();
            let ty = &field.ty;

            let mut is_skip = false;

            let attr = field
                .attrs
                .iter()
                .find(|attr| attr.path().is_ident("optional"));

            if let Some(attr) = attr {
                let _ = attr.parse_nested_meta(|meta| {
                    if meta.path.is_ident("skip") {
                        is_skip = true;
                    } else if meta.path.is_ident("rename") {
                        let renamed_ident = meta.value()?.parse::<syn::Ident>()?;
                        ident = Some(renamed_ident);
                    }
                    Ok(())
                });
            }

            if is_skip || is_option(ty) {
                return quote::quote!(#ident: #ty);
            }

            quote::quote!(#ident: Option<#ty>)
        })
        .collect();

    quote::quote! {
        struct #optional_struct_ident {
            #(#fields,)*
        }
    }
    .into()
}

fn is_option(ty: &Type) -> bool {
    let Type::Path(path) = ty else {
        return false;
    };
    let path = &path.path;
    path.segments
        .last()
        .map(|segment| segment.ident == "Option")
        .unwrap_or(false)
}

总结

通过上面的示例, 可以掌握基本的派生宏的编写方式

你可以自己定义宏, 然后生成你想要的代码, 可以考虑下假设有生命周期和泛型的时候, 该如何编写

本文是宏魔法的第一篇, 后面会继续介绍 rust 其他的宏和它编写方式