如何在 Rust 中实现一个自定义宏

一、前言

最近在用 Rust 刷算法题的过程中遇到了几道二叉树相关的问题，题目给的输入基本都是二叉树层序遍历的结果（如下图所示），在本地测试的话需要根据其自行构造二叉树。

这部分的具体代码实现很常见，因此不再详细介绍。最后的形式如下面代码所示，其中传入参数的类型为Vec<Option<T>>，即二叉树层序遍历的结果。

rust 复制代码

TreeNode::from(vec![Some(1), Some(5), Some(3), None, Some(4), Some(10), Some(6), Some(9), Some(2)])

可以看到参数的形式比输入示例复杂了不少，每次在本地测试都需要手动输入测试用例，略显麻烦。本着省事的原则（~~懒是第一推动力~~），就想通过某些方法实现可以直接根据测试用例的输入构造二叉树。

但是测试用例给的输入是 JSON 格式的数据，Rust 本身不支持类似的语法，虽然可以通过 serde 等库来对 JSON 进行反序列化，但是其难免会引入运行时开销。而宏编程 可以在不引入运行时开销 的前提下完成对自定义语法的扩展，完美符合要求。

rust 复制代码

// 测试用例的输入
// [1, 5, 3, null, 4, 10, 6, 9, 2]

// 使用 serde 进行 JSON 反序列化
let vals: Vec<Option<i32>> = serde_json::from_str("[1, 5, 3, null, 4, 10, 6, 9, 2]").unwrap();
let root = TreeNode::from(vals);

// 使用类函数宏
let root = tree![1, 5, 3, null, 4, 10, 6, 9, 2];

二、Rust 中的宏编程

参考 Rust 圣经中关于宏的介绍：宏 - Rust 程序设计语言简体中文版 (kaisery.github.io)

正如上面所说，宏编程可以在不引入运行时开销的同时创建自定义语法扩展，其本质上是一种在编译阶段通过代码生成代码的方式，能帮助简化代码和减少大量重复的代码，提升代码可读性和可维护性。

在 Rust 中，宏有声明宏 （Declarative Macros）和过程宏 （Procedural Macros）两种类型，其中声明宏通过模式匹配进行定义，而过程宏通过函数进行定义。而过程宏又包含派生宏 （Derive Macros）、类属性宏 （Attribute-like Macros）和类函数宏（Function-like Macros）三种。

1. 声明宏

声明宏是 Rust 中最常用的宏类型，比如常见的println!和vec!都属于声明宏。声明宏通过macro_rules!定义，其接收一个表达式 ，通过对表达式的结果进行模式匹配 来执行对应代码。一个vec!宏定义的简化版本如下面代码所示：

rust 复制代码

#[macro_export]
macro_rules! vec {
    ( $( $x:expr ),* ) => {
        {
            let mut temp_vec = Vec::new();
            $(
                temp_vec.push($x);
            )*
            temp_vec
        }
    };
}

2. 派生宏和类属性宏

派生宏通常用于为结构体生成 trait 的默认实现 ，比如常见的#[derive(Debug)]等。在下面这个简单的例子中，用户可以使用#[derive(Hello)]注解它们的类型来得到hello函数的默认实现。而对于派生宏的定义，需要实现一个函数，将结构体代码作为参数输入，其类型为TokenStream，并将生成的扩展代码作为结果返回，其类型同样为TokenStream。

rust 复制代码

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

#[proc_macro_derive(Hello)]
pub fn hello_macro_derive(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput); // 将输入 TokenStream 转换为抽象语法树
    let ident = ast.ident; // 获取结构体的名称
    let output = quote! {
        impl Hello for #ident {
            pub fn hello() {
                println!("Hello world!");
            }
        }
    }; // 使用 quote! 宏扩展代码
    output.into()
}

rust 复制代码

trait Hello {
    fn hello();
}

#[derive(Hello)]
struct Pancakes;

// impl Hello for Pancakes {
//     pub fn hello() {
//         println!("Hello world!");
//     }
// }

fn main() {
    Pancakes::hello();
}

类属性宏与派生宏类似，不同的是派生宏只能用于结构体和枚举 ，而类属性宏还可以用于其它的项，比如函数。类属性宏的定义如下面代码所示：

rust 复制代码

#[proc_macro_attribute]
pub fn attribute_like_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
    ...
}

rust 复制代码

#[attribute_like_macro(...)]
fn test() {
    ...
}

3. 类函数宏

类函数宏和声明宏一样在使用上都与函数调用 类似，但由于使用函数定义的方法，其比声明宏更为灵活。类函数宏的定义如下面代码所示，获取括号中的代码，并返回希望生成的代码：

rust 复制代码

#[proc_macro]
pub fn func_like_macro(input: TokenStream) -> TokenStream {
    ...
}

rust 复制代码

func_like_macro!(...);

三、自定义类函数宏

正如上述关于 Rust 中不同类型的宏的介绍，使用类函数宏来解决前言中提到的问题是最为合适的，因此下面介绍如何实现这样一个自定义的类函数宏。

1. 语法分析

抽象语法树（Abstract Syntax Tree, AST）是源代码语法结构的一种抽象表示，以树状数据结构表示编程语言的语法结构。虽然宏编程可以实现自定义语法扩展，但最终还是需要解析成 AST，因此下面先进行语法分析。

这里推荐一个网站：AST explorer，可以将输入的代码解析成 AST，可以方便我们进行语法分析。

这里重点看上图中标记的TokenStream部分，这便是类函数宏定义的输入，可以看到列表中的内容被解析成了三种不同类型的 token：Literal （字面量）、 Ident （标识符）和 Punct（标点）。

2. 代码实现

通过语法分析，结合宏扩展的目标，我们发现只需要在遍历TokenStream的过程中对字面量和除null之外的标识符使用Some进行包裹，而将null转换为None。具体代码如下：

rust 复制代码

#[proc_macro]
pub fn tree(input: TokenStream) -> TokenStream {
    let mut vals = Vec::new();
    let mut is_comma = false;
    for token in input {
        if let TokenTree::Punct(_) = &token {
            if is_comma {
                vals.push(String::from("None")); // 兼容连续出现两个逗号的情况
            }
        } else {
            let token_str = token.to_string();
            let val = if token_str == "None" || token_str == "null" {
                String::from("None")
            } else {
                format!("Some({token})")
            };
            vals.push(val);
            is_comma = false;
        }
    }
    let output = format!("TreeNode::from(vec![{}])", vals.join(", ")); // 拼接生成的代码
    output.parse().unwrap()
}