Rust 属性宏从入门到发布

简介

本篇主要介绍如何实现一个可以打印函数执行时间的 Rust 属性宏并且发布到 crates.io(类似于 npmjs.com) 中，用法如下：

use core::time;
use std::thread;

use puff_bench_macro::log_bench;

#[log_bench]
fn test_my_macro() {
    for num in 0..3 {
        thread::sleep(time::Duration::from_secs(1));
        println!("{}", num);
    }
}

fn main() {
    test_my_macro();
}

上面的打印结果：

进入函数: test_my_macro
0
1
2
离开函数: test_my_macro (耗时 3025 ms)

log_bench 是一个属性宏，它可以把当前函数的执行时间打印出来。最后的打印结果如下

你可以理解属性宏类似于 Java 的注解或者类似于 TypeScript 的装饰器函数，但是也仅仅只是类似，它们底层实现原理不一样。Rust 是直接在编译时去做代码注入的相关工作的，接下来我们就来看下如何实现这么一个宏吧，不过在开始编码之前，要先看看一些前置的基础知识。

前置知识

什么是 Rust 属性宏？

Rust 属性宏是一种元编程工具，它允许开发者在 Rust 代码中使用自定义的属性（attribute）来影响编译时的代码生成。也就是这个工作是在编译阶段去做的，可以通过宏去添加/修改调用这个属性宏的代码，它的使用场景包括创建领域特定语言（DSL）、自定义的代码生成器等，以满足特定需求。

一些基础知识点：

• TokenStream： 宏的输入和输出通常是 TokenStream，它表示一系列的 token。这些 token 可以包括标识符、关键字、操作符等。宏通过操作这些 token 来生成、修改或分析代码
• 属性（Attributes）： 在 Rust 中，属性是以 #[...] 格式的注解，附加在项（item，如函数、结构体、枚举等）上，属性宏可以拿到调用它的比如函数/结构体等的信息，然后做一些相关的操作

常用的开发宏的库:

• **syn：**这是一个用于解析 Rust 代码的库，它允许开发者将代码解析为抽象语法树（AST）。在属性宏中，通常需要使用 syn 来分析输入的 TokenStream。
• quote： 这是一个用于生成 Rust 代码的库，它允许开发者使用类似于模板的语法来创建代码片段，然后将这些片段嵌入到生成的代码中。在属性宏中，通常需要使用 quote 来构建和生成输出的 TokenStream。
• **proc-macro2：**这是 syn 的一个分支，提供了一些额外的工具和功能，用于更灵活地处理宏。在属性宏的开发中，proc-macro2 通常会和 syn 一起使用

代码实现

首先创建一个库项目

cargo new puff_bench_macro --lib

然后我们到 cargo.toml 中加点东西

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

[lib]
proc-macro = true # 告诉编译器当前 crate 包含过程宏。Rust 编译器将使用这个标记来正确处理过程宏的编译和链接过程

[[bin]]
name = "main"
path = "bin/main.rs"

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

[dev-dependencies]
puff_bench_macro = {path = "."}

然后我们就到 src/lib.rs 中编写代码

// 引入 proc_macro crate，用于创建过程宏
extern crate proc_macro;
use proc_macro::TokenStream;
// 引入 quote crate，用于生成代码片段
use quote::quote;
// 引入 syn crate，用于解析 Rust 代码
use syn::parse_macro_input;

// 定义 log_bench 属性宏
#[proc_macro_attribute]
pub fn log_bench(_: TokenStream, item: TokenStream) -> TokenStream {
    // 解析传入的函数定义
    let input_fn = parse_macro_input!(item as syn::ItemFn);

    // 获取函数名和函数体
    let fn_name = &input_fn.sig.ident;
    let fn_block = &input_fn.block;

    // 构建新的函数定义，包含性能日志输出
    let expanded = quote! {
        fn #fn_name() {
            // 获取函数执行开始时间点
            let start = std::time::Instant::now();
            println!("进入函数: {}", stringify!(#fn_name));
            // 执行原函数体
            #fn_block
            println!("离开函数: {} (耗时 {} ms)", stringify!(#fn_name), start.elapsed().as_millis());
        }
    };

    // 将生成的代码片段转换为 TokenStream，以便返回
    TokenStream::from(expanded)
}

当我们执行这段代码的时候,其实就是去获取它的函数信息:

let input_fn = parse_macro_input!(item as syn::ItemFn);

ItemFn 的类型定义如下:

pub struct ItemFn {
    pub attrs: Vec<Attribute>, // 包含函数所有属性
    pub vis: Visibility, // 调用函数的可见性,比如 pub pub(crate)
    pub sig: Signature, // 包含了函数的签名信息，如函数名、参数、返回类型等
    pub block: Box<Block>, // 函数体的代码块，即函数的实际执行部分1
}

然后我们就可以在后面利用这些信息重新组合成一个新的代码片段并返回。

到了这里我们的宏就已经实现完成了, 然后编写一段测试代码来看看吧

测试

在根目录下创建 bin/main.rs 文件, 然后在它里面去编写测试代码:

use core::time;
use std::thread;

use puff_bench_macro::log_bench;

#[log_bench]
fn test_my_macro() {
    for num in 0..3 {
        thread::sleep(time::Duration::from_secs(1));
        println!("{}", num);
    }
}

fn main() {
    test_my_macro();
}

此时你就可以在命令行运行 cargo run 命令，就能够看到对应的打印信息了

发布

最后你可以把这个包发布到 crates.io 上面，让所有人都能用到这个，执行两个命令即可

cargo login

复制你的 token 到命令行就能登陆成功

cargo publish

大功告成