Rust 学习之自定义派生宏-访问和操作TokenStream(一)-道・自然之境

在自定义派生宏中访问和操作 TokenStream 的核心思路是：先解析为结构化的语法树（用 syn 库），再基于语法树信息生成新的代码（用 quote 库），最后转换回 TokenStream。直接操作原始 TokenStream（如单个令牌）很少见，因为它本质是一串无序的 “语法碎片”，而 syn 库提供了结构化的访问方式，更易于处理。

一、核心流程：从 `TokenStream` 到新代码

操作 TokenStream 的完整流程可以概括为三步：解析输入 TokenStream → 处理语法树信息 → 生成新的 TokenStream

步骤 1：解析输入 `TokenStream`（用 `syn` 库）

TokenStream 是一串原始令牌（如 struct、{、id 等），无法直接获取 “结构体名称”“字段列表” 等语义信息。必须用 syn 库将其解析为语法树结构体（如 DeriveInput），才能方便地访问类型信息。

关键函数：syn::parse_macro_input!

这个宏会将 TokenStream 解析为指定的语法树类型（如 DeriveInput，表示一个可被派生的类型，包括结构体、枚举等），并自动处理解析错误（返回友好的编译错误）。

示例：解析结构体的 TokenStream：

use proc_macro::TokenStream;
use syn::{DeriveInput, Fields, Ident}; // DeriveInput 是 syn 定义的“可派生类型”语法树

#[proc_macro_derive(MyMacro)]
pub fn my_macro(input: TokenStream) -> TokenStream {
    // 1. 解析 input 为 DeriveInput（包含结构体/枚举的全部信息）
    let input = syn::parse_macro_input!(input as DeriveInput);

    // 2. 从 DeriveInput 中提取关键信息
    let struct_name = input.ident; // 结构体名称（如 User）
    let fields = match input.data {
        // 只处理结构体（忽略枚举等其他类型）
        syn::Data::Struct(s) => s.fields,
        _ => panic!("MyMacro 只支持结构体！"),
    };

    // 3. 进一步提取字段信息（以“带名称的字段”为例，如 struct A { x: i32 }）
    let field_names: Vec<Ident> = match fields {
        Fields::Named(named) => named
            .named
            .iter()
            .map(|field| field.ident.clone().unwrap()) // 字段名（如 x）
            .collect(),
        _ => panic!("MyMacro 只支持带名称的字段（如 { x: i32 }）！"),
    };

    // 后续步骤：基于 struct_name 和 field_names 生成代码...
}

步骤 2：处理语法树信息（自定义逻辑）

解析得到语法树后，就可以根据业务需求处理信息了。例如：

记录结构体名称，用于生成 impl 结构体名；
遍历字段列表，生成字段相关的代码（如打印字段名、生成访问方法等）。

接上面的示例，假设我们要收集所有字段名，用于后续生成代码：

// 接上面的代码
println!("解析到结构体：{}", struct_name); // 打印调试信息（编译期可见）
println!("字段列表：{:?}", field_names); // 输出：[id, name]（假设字段是 id 和 name）

步骤 3：生成新的 `TokenStream`（用 `quote` 库）

处理完信息后，需要用 quote! 宏生成新的 Rust 代码，并转换为 TokenStream 返回给编译器。quote! 支持嵌入变量（用 # 符号），自动将语法树结构转换为对应的令牌流。

示例：生成一个打印所有字段名的方法：

use quote::quote;

// 接上面的代码
// 4. 用 quote! 生成代码，嵌入 struct_name 和 field_names
let generated_code = quote! {
    // 为结构体实现一个打印字段名的方法
    impl #struct_name {
        pub fn print_fields() {
            let fields = vec![#(#field_names.to_string()),*]; // 展开字段名列表
            println!("结构体 {} 的字段：{:?}", stringify!(#struct_name), fields);
        }
    }
};

// 5. 将生成的代码转换为 TokenStream 并返回
TokenStream::from(generated_code)

#struct_name：嵌入结构体名称（如 User）；
#(#field_names.to_string()),*：循环展开字段名列表（如 id.to_string(), name.to_string()）；
quote! 的返回值是 proc_macro2::TokenStream，需通过 TokenStream::from() 转换为 proc_macro::TokenStream。

二、完整示例：实现一个简单的派生宏

将上述步骤整合，实现一个 #[derive(PrintFields)] 宏，为结构体生成 print_fields 方法：

// src/lib.rs
use proc_macro::TokenStream;
use quote::quote;
use syn::{DeriveInput, Fields, Ident};

#[proc_macro_derive(PrintFields)]
pub fn derive_print_fields(input: TokenStream) -> TokenStream {
    // 步骤 1：解析输入 TokenStream 为语法树
    let input = syn::parse_macro_input!(input as DeriveInput);
    let struct_name = input.ident;

    // 步骤 2：提取字段信息
    let field_names: Vec<Ident> = match input.data {
        syn::Data::Struct(s) => match s.fields {
            Fields::Named(named) => named
                .named
                .iter()
                .map(|field| field.ident.clone().unwrap())
                .collect(),
            _ => panic!("PrintFields 只支持带名称的字段（如 { x: i32 }）！"),
        },
        _ => panic!("PrintFields 只支持结构体！"),
    };

    // 步骤 3：生成代码并返回 TokenStream
    let generated_code = quote! {
        impl #struct_name {
            pub fn print_fields() {
                let fields = vec![#(#field_names.to_string()),*];
                println!("结构体 {} 的字段：{:?}", stringify!(#struct_name), fields);
            }
        }
    };

    TokenStream::from(generated_code)
}

三、测试宏：使用生成的代码

在另一个文件中使用 #[derive(PrintFields)]：

// src/main.rs
use my_macro::PrintFields;

#[derive(PrintFields)]
struct User {
    id: u32,
    name: String,
    age: u8,
}

fn main() {
    User::print_fields(); 
    // 输出：结构体 User 的字段：["id", "name", "age"]
}