Rust 宏系列教程—自定义属性宏

在Rust中，属性宏（Attribute Macros）是一种强大的元编程工具。它们允许开发者通过自定义属性（以#[...]的形式）来扩展Rust语言的语法，从而在编译时对代码进行转换。属性宏在编译阶段起作用，能够读取和修改Rust抽象语法树（AST），以生成新的代码或者修改现有代码的行为。

本文带你从头自定义属性宏，用于自定义评估函数执行占用时间。

创建项目工程

首先创建主体项目 foo2, 然后进入该项目目录中，再创建macro_lib 项目，我们在子项目中实现属性宏代码。

# 创建主体工程
cargo new foo2 

cd foo2
cargo new macro_lib --lib

项目树结构大致如下：

├── src
│   └── main.rs
├── macro_lib
│   ├── src
│   │   ├── log_duration.rs
│   │   └── lib.rs
│   ├── Cargo.toml
│   └── Cargo.lock
├── Cargo.toml
└── Cargo.lock

在主项目的cargo.toml中增加子项目依赖：

[dependencies]
macro_lib = { path = "./macro_lib" }

声明属性宏

我们在子项目的lib.rs中定义函数，并使用宏进行标识，告诉编译器该函数是宏声明：

// macro_lib/src/lib.rs

#[proc_macro_attribute]
pub fn log_duration(args: TokenStream, item: TokenStream) -> TokenStream {
    log_duration_impl(args, item)
}

对于属性宏定义，函数名很重要，因为名称就是属性宏的名称。正如你所看到的，这需要两个不同的参数。第一个是传递给属性宏的参数，第二个是属性宏的目标。

为了实现于声明分离，我们定义新的文件实现log_duration_impl函数：

touch src/log_duration.rs

实现log_duration 属性宏

我将首先给你完整的实现，然后我将分解到目前为止我还没有使用的部分：

// macro_lib/src/log_duration.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn};

pub(crate) fn log_duration_impl(_args: TokenStream, input: TokenStream) -> TokenStream {
    // 解析输入作为ItemFn类型，它是syn 提供的表示函数类型
    let input = parse_macro_input!(input as ItemFn);

    let ItemFn {
        // 函数签名
        sig,
        // 函数可见性标识
        vis,
        // 函数体
        block,
        // 其他属性
        attrs,
    } = input;

    // 抽取函数体语句
    let statements = block.stmts;

    // 存储函数名称标识符，用于日志记录
    let function_identifier = sig.ident.clone();

    // 使用解析输入重构函数，然后输出
    quote!(
        // 在该函数上重复其他所有属性（保持不变）
        #(#attrs)*
        // 重构函数声明
        #vis #sig {
            // 记录开始时间
            let __start = std::time::Instant::now();

            // 创建新的块，即函数的主体部分，存储返回值作为变量，后续可返回给父函数
            let __result = {
                #(#statements)*
            };

            // 记录函数执行时间
            println!("{} took {}μs", stringify!(#function_identifier), __start.elapsed().as_micros());

            // 返回结果
            return __result;
        }
    )
    .into()
}

你之前可能没有看到的可能是：通过将输入解析为ItemFn而获得的签名和块字段。Sig包含函数的整个签名，而block包含函数的整个主体。这就是为什么，通过使用下面的代码，我们基本上可以重建未修改的函数：

// 在宏里面重建未修改fn

#vis #sig #block

在本例中，你希望修改函数体，这就是为什么要创建封装原始函数块的新块。我们现在以及实现了属性宏，这里补充下lib.rs代码及依赖：

cargo.toml

[package]
name = "macro_lib"
version = "0.1.0"
edition = "2021"

[lib]
name="macro_lib"
proc-macro = true
path = "src/lib.rs"

[dependencies]
quote = "1.0.37"
syn = {version ="2.0.87", features=["full"] }

proc-macro = true 这个是标识该项目是过程宏项目。当然还需要引用 quote 和 syn

lib.rs完整代码

mod log_duration;

extern crate proc_macro;

use proc_macro::TokenStream;
use log_duration::log_duration_impl;

#[proc_macro_attribute]
pub fn log_duration(args: TokenStream, item: TokenStream) -> TokenStream {
    log_duration_impl(args, item)
}

应该具体实现在log_duration文件中，因此需要声明mod，接着是引入log_duration_impl。下面我们看如何使用该宏。

使用属性宏

// main.rs

use macro_lib::log_duration;

#[log_duration]
#[must_use]
fn function_to_benchmark() -> u16 {
    let mut counter = 0;
    for _ in 0..u16::MAX {
        counter += 1;
    }

    counter
}

fn main() {
    println!("{}", function_to_benchmark());
}

输出结果：

function_to_benchmark took 498μs
65535

总结

本文完整实现了简单的属性宏，并采用声明与实现分离方式实现。如果你之前以及阅读本系列文章，应该不会感动很难懂。下面我们会继续更复杂的属性宏实现，来吧，一起学习rust！