Rust 过程宏入门

之前写过几篇乱七八糟的 Rust 入门的流水账 1 2 3,并没有仔细介绍过 Rust 的宏系统,稍微补一下。

Rust 强大的宏系统是图灵完备的(对, C/C++ 的宏也是图灵完备的,但是就是用着膈应啊),甚至可以控制编译器的行为,最近进入 stable 的过程宏 (Procedural Macro) 相当于是在代码中给编译器打插件。 Rust 的宏能干的事情有:代码生成(基础功能)、获取自定义数据类型结构并生成代码(序列化和反序列化)、获取函数实现并修改代码(类似、甚于 Python 的 decorator )。

举例来说, Rust 的宏允许你写出这样的代码:

let map = hashmap! {
    "Java" => 1,
    "C" => 2,
    "Python" => 3,
};

或者这样的代码:

#[derive(Serialize, Deserialize)]
struct Struct {
    #[serde(rename = "i32")]
    i: i32,
}

甚至是这样的代码:

use rocket_contrib::json::Json;
use std::collections::HashMap;

#[post("/register", format="application/json", data="<data>")]
fn register(data: Json<HashMap<String, String>>) -> String {
    format!("{:#?}", data.into_inner())
}

上面所述的这些宏的使用方式都是已经有项目实现并在使用的,如序列化和反序列化库 serde ,如使用 Flask 路由风格的、占据 Rust 1/4 HTTP 服务库江山的 rocket 等等。

因为感觉互联网上关于过程宏的中文资料比较少,而且由于 Rust 换代极快,包括很多外文资料、链接都已经失效,所以找老外的博文加上自己读了一下几个库的源码稍微学习了一波,总结一下。

官方资料:

macro_rules!

还是从最基本的宏定义开始。在 Rust 学习过程中首先接触的一般都是 println! 打印宏(输出 Hello World )。这类调用方式极其类似普通函数的宏的定义就是 macro_rules! 的工作。

关于这种宏的定义技巧学习推荐资料 The Little Book of Rust Macros ,里面介绍了宏的各种性质,从基本的定义到实践中的各种技巧,例子也十分详尽。是 Rust 宏的必考知识点手册。

因为这本书有中文翻译《Rust宏小册 中文版》,所以具体内容我也不再介绍,就当是娱乐,写个帮助宏。

macro_rules! count_tt {
    ()      => (0usize);
    ($e:tt) => (1usize);
    ($($s:tt $t:tt)*) => (
        count_tt!($($t)*) << 1 | 0
    );
    ($e:tt $($s:tt $t:tt)*) => (
        count_tt!($($t)*) << 1 | 1
    );
}

这段代码从 Rust 1.2.0 开始均可使用。这个宏的作用是计算 tt 的个数。用的二分法,所以按 Rust 默认 32 层宏展开的限制,可以计算 4294967296 个标记,而且结果也是编译期常量。这个比 The Little Book of Rust Macros 中给出的几种方法好用多了,虽然……我也觉得不会有在源码里生成出超过一万个标记的情况,而且巨慢……

过程宏

如果说 macro_rules! 是使用递归和模式匹配、字符串替换的函数式风格定义宏的话,过程宏 #[proc_macro] 则是更贴近 Rust 本身过程式语法的定义方法。

Rust 的过程宏支持扩展的语法更加丰富,它可以定义

  • 拟函数语法,和 macro_rules! 定义的宏调用规则类似;
  • 自定义派生,如 #[derive(Debug)]
  • 拟属性语法,如 #[route(GET, path="/")]

因为过程宏需要手动处理标记并输出新的一段符合 Rust 语法的标记,为了方便解析、生成,一般会引入 synquote 库。这两个库是同一个作者编写的,前者将符合 Rust 语法的 proc_macro::TokenStream 解析成容易访问的结构,后者将 Rust 代码块重新输出为标记流。有这两个库的帮助,编写过程宏就变得非常容易了。

某种意义上来说,本文可能会沦为 syn 和 quote 库的使用教程(笑)。

拟函数语法

如果看完了 The Little Book of Rust Macros 应该可以学到一系列 macro_rules! 递归处理标记的技巧。但是,有些情况下使用 macro_rules! 来定义宏并不适合,甚至可以说无法胜任。主要是这种宏的定义还有递归深度限制和一系列语法规则的限制,还要算上这种宏只能通过已定义的规则匹配和生成,它无法计算!——就是说你无法从一个 0 生成一个 1 ,从 1 生成 2 等等, macro_rules! 本质只是字符串替换和模式匹配,不会进行求值。因此这种情况下程序员只能对每个数字进行手动匹配。这就导致某些情况下并不适用 macro_rules! 。而相对来说,使用过程式则更加自由、便于调试。

说这么多不如来举个简单的例子。

由于 proc-macro 必须单独作为一个 crate 提供,因此这里不得不 cargo new 两次。此处用来示范的 crate 名叫 hw , proc-macro crate 名为 pm ,则初始化为:

cargo new --bin hw
cd hw
cargo new --lib pm

于是接下来的描述均以 git 仓库根目录——也就是 hw ——为当前工作目录了。

因为 hw 引用了 pm ,所以需要在 Cargo.toml 中添加一条 dependencies

[dependencies]
pm = { path = "pm" }

由于 pm 包需要使用 proc-macro 特性,所以需要在 pm/Cargo.toml 中添加两行开启 proc-macro :

[lib]
proc-macro = true

最简单的,写一个接受任意参数都输出 Hello, World! 的宏,编辑 pm/src/lib.rs 为

extern crate proc_macro;

use proc_macro::TokenStream;

#[proc_macro]
pub fn hw(_: TokenStream) -> TokenStream {
    r#"println!("Hello, World!");"#.parse().unwrap()
}

最后生成的宏名是和函数名相同的,也就是能用的宏就是 hw!

具体语法不再解释。由于 proc macro 是以 crate 的形式提供的,因此需要引入 proc_macro 包, TokenStream 是一个结构体,它是标记流,作为函数参数的 TokenStream 是在函数作为宏被调用的时候的传入的参数,作为返回值的 TokenStream 是函数处理完输入后返回给编译器的标记流。也就是说,在编译器处理到 hw! 宏的时候会将 hw! 所有参数打包为 TokenStream 交给上面定义的函数执行,然后用返回的 TokenStream 替换掉原先宏所占据的源码位置。

然后可以试着在 src/main.rs 中调用 hw!() 宏,可以发现程序执行的时候输出了一行 Hello, World! ,这事实上就是上面代码中一串 str parse 成 TokenStream 替换 hw!() 的程序的运行结果。由于 TokenStream 和 String 以及 str 都有互转接口,因此在过程宏中手写语法分析实现其它语言的编译器还是有可能的。

当然,就上面那个实例函数而言,一点用处都没有。接下来用 proc macro 重写一个 count_tt! 函数。 TokenStream 似乎没有直接计算长度的函数,不过实现了 IntoIterator 所以可以迭代一遍统计个数:

#[proc_macro]
pub fn count_tt(ts: TokenStream) -> TokenStream {
    ts.into_iter().count().to_string().parse().unwrap()
}

虽然看起来很蠢,不过速度居然比一开始用 macro_rules! 实现的要快,还算是一个不错的选择。

自定义派生

最常见的过程宏应用就是在定义 traits 的时候自定义派生类了。在自定义一个新的类的时候,往往会选择直接 #[derive(Debug)] 以便 print! 宏直接输出。利用过程宏的 proc_macro_derive 属性,现在普通程序员也可以做到为自己的 trait 撰写自定义派生。

好了, syn 和 quote 教程正式开始。目前我是用的是写文时最新的 syn 和 quote ,也就是:

[dependencies]
syn = "0.15"
quote = "0.6"

再次介绍一下 syn 和 quote ,这两个库作者都是 @dtolnay ,他同时也是 serde 的开发者之一。 serde 是一个对 Rust 结构体序列化和反序列化的库,支持许多数据格式,默认支持 JSON ,也支持 YAML 、 MessagePack 、 x-www-form-urlencoded 、 Pickle 等等一系列序列化格式。而 serde 本身也依赖于 syn 和 quote 两个库——不如说,这两个库一开始就是为 serde 而生的。

而如上文所说的, syn 库负责将 Rust 代码的 TokenStream 解析为更容易被程序员访问的结构块,而 quote 库负责从 Rust 代码重新生成 TokenStream ,以便返回给编译器。

实际的例子中,我们需要实现一个 trait ,只有一个函数 fn hw(&self) -> () ,默认情况下它获取结构体的名称并将其打印出来,简单地写在 src/main.rs 中:

trait HelloWorld {
    fn hw(&self) -> ();
}

我们希望程序员可以直接通过 #[derive(HelloWorld)] 属性来为类自动实现这个 trait 。这时候 proc_macro_derive 就出场了,它可以为 derive 属性实现新的 trait ,在 pm/src/lib.rs 中我们如此写:

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;

#[proc_macro_derive(HelloWorld)]
pub fn hw(ts: TokenStream) -> TokenStream {
    let ast: syn::DeriveInput = syn::parse(ts).unwrap();
    let name = &ast.ident;
    let gen = quote! {
        impl HelloWorld for #name {
            fn hw(&self) -> () {
                println!("Hello, {}!", stringify!(#name));
            }
        }
    };
    gen.into()
}

我想不需要一行一行解释吧,大概流程是使用 syn 库 parse TokenStream 为 ast: syn::DeriveInput ,从 ast.ident 获取到类名,然后使用 quote 库的特殊宏将类名嵌入到代码中重新生成一个 proc_macro2::TokenStream , Into trait 实现了到 proc_macro::TokenStream 的转换。

看看,非常好用。不过把玩上述代码的过程中很容易遇到一个问题,那就是泛型怎么整?泛型 impl 块比普通的要多了那么几个泛型的声明。

也不需要担心, syn 库为我们提供了一站式的接口 split_for_impl 。

我们先考虑给一个泛型类实现 trait 会写成什么样。这是 Rust 入门级别的问题,此处略过,大概长这样: impl<T> Trait for Struct<T: Constraint> where WhereClause 。所以我们的代码就——

#[proc_macro_derive(HelloWorld)]
pub fn hw(ts: TokenStream) -> TokenStream {
    let ast: syn::DeriveInput = syn::parse(ts).unwrap();
    let name = &ast.ident;
    let (im, ty, wh) = ast.generics.split_for_impl();
    let gen = quote! {
        impl #im HelloWorld for #name #ty #wh {
            fn hw(&self) -> () {
                println!("Hello, {}!", stringify!(#name))
            }
        }
    };
    gen.into()
}

即可。正如所见, #im 是跟在 impl 之后的泛型声明, #ty 是跟着类型后的给类的泛型定义, #wh 则是最后的 where 限制。当任意一个不存在的时候,对应的变量在 quote! 中就是空。换句话说,这段代码也可以直接处理不带泛型的普通结构体、类(毕竟一站式接口)。

不过还有一个小问题:输出的时候并不会带上泛型类参数。当然这个问题还是可以解决的,如果不考虑格式化和对每个实例都有唯一的 hw 输出,且不考虑格式化的话,将上文的 #ty 拼接到 #name 后就可以了。

但是如果需要格式化就是一个比较大的问题了,逃不开递归对 T 求类型名和字符串处理,因此这里我建议使用 std::intrinsics::type_name 。🤣🤣🤣

当然如果你不想在代码中引入新的不稳定的 feature 和 unsafe 的话,可以考虑使用 typename 包,可以去它的 GitHub 鉴赏一下其源码,看看别人是怎么写的(这个包也使用 proc macro 实现的)。

syn quote 包简介

Rust 的包管理非常友善,统一由 crates.io 进行注册唯一标识。再加上 rustdoc 这个工具可以利用源码注释生成文档,这也使得包文档的整理更加容易。基本上学习一个简单的库的使用只需要啃文档就好了。这其中的 syn 和 quote 就是两个非常容易上手的库,但是如果需要深度挖掘它们的潜力的话,还是好好啃一遍文档比较好。

在上文中简单地使用了一遍 syn 和 quote ,这里打算稍微深入一点点讲解一下基本用法,方便调试等。

首先是 syn ,它在包中定义了大量的 Structs 和 Enums ,对应的是 Rust 源码中的各种元素。在 derive 中我们只使用到了 syn::DeriveInput 相关的各种类型,如 ast.ident 成员对应的 Ident 类型, ast.generics 成员对应的 Generics 类型等。后文中主要用的是 syn::Item中相关的一系列类型。

仔细研究每个类,会发现许多都 impl 了一个来自 quote 库的 trait ,即 quote::ToToken 。事实上, quote 库正是通过这个 trait 将在 quote!块中出现的插值变量转换为标记流的。

ToTokens trait 有两个函数,分别实现了 to_tokensinto_token_stream ,通过它们可以将 syn 中的一部分类转化为 proc_macro2::TokenStream ,然后可以通过 to_string 再变回我们熟悉的字符串——就可以在编译期调试输出了(误)。

除此之外, quote! 块中还支持重复插值,也就是和 macro_rules! 中 $()* 类似的操作,只需要将 $ 替换为 # 即可。只要实现了 IntoIterator 这个 trait 的类型,包括 std::vec::Vec 都可以在 quote! 块中使用重复展开(当然 Iterator 解引用的类需要 impl ToTokens 呀) 。

接下来事情就简单多了,当我们遇到 struct 就去查成员,当我们遇到 enum 就去匹配元素。至于每个类型啥意思嘛……读 (cai) 文 (ming) 档 (zi) 呗。

拟属性语法

Rust 过程宏的拟属性语法相比于 Python 装饰器的强大之处在于,它可以通过修改 Token 流对函数内部、结构体内部进行修改。当然一般来说这种操作过于复杂不会有人这么做。相比之下简单的操作可以有获取函数参数、获取函数返回值、获取结构体成员等等一系列操作。

拟属性语法和自定义派生的区别是,自定义派生通过获取标记流生成一串新的标记流附加到其后,它不会更改原代码;而拟属性语法则是直接在标记流上修改,返回值就是修改后的标记流。

由于属性可以标记到几乎任何代码结构上,因此 syn 的使用方法和自定义派生中略有不同,不应当再用 syn::DeriveInput 来作为 parse 结果,而是一个 Enum syn::Item

不过想要使用该类型,首先需要开启 syn 的 “full” feature ,于是在 pm/Cargo.toml 改为

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

[lib]
proc-macro = true

举个简单的例子,我希望属性只作为函数的标记,在作为函数标记的情况下什么都不做(也就是单纯地返回该函数),其它情况下报错:

#[proc_macro_attribute]
pub fn hw(_head: TokenStream, body: TokenStream) -> TokenStream {
    match syn::parse::<syn::Item>(body).unwrap() {
        syn::Item::Fn(ref func) => {
            use quote::ToTokens;
            func.into_token_stream().into()
        }
        _ => panic!("Only fn is allowed!"),
    }
}

代码很好理解,可以写一个小程序验证一下,只有在作为函数属性的时候才不会报错,其他时候都会提示 help: message: Only fn is allowed! ,也就是 panic! 的内容。在过程宏中 panic! 会使编译器抛出一个编译错误。在 unstable 或者 nightly 中还可以使用 proc_macro::Diagnostic 类产生更加明确的提示,如指定一段代码区间,使编译器高亮以提示此处有误,等等。不过现在实现起来略有些麻烦,这里就懒得写了。

有了上面的基础,就很容易可以实现 Python 的装饰器类似的功能了,只要在 match 到函数之后, 获取到函数的各项属性并在 quote! 中重新定义即可。我打算写一个宏属性使得被宏属性标记的函数会在开始和结尾输出 begin 和 end 。那么就有很简单粗暴的实现——

#[proc_macro_attribute]
pub fn hw(_head: TokenStream, body: TokenStream) -> TokenStream {
    match syn::parse::<syn::Item>(body).unwrap() {
        syn::Item::Fn(ref func) => {
            let syn::ItemFn {
                attrs, vis, constness, unsafety,
                asyncness, abi, ident, decl, block,
            } = func;
            use std::ops::Deref;
            let syn::FnDecl {
                generics, inputs, variadic, output, ..
            } = decl.deref();
            let gen = quote! {
                #(#attrs)* #vis #constness #unsafety #asyncness #abi
                fn #ident #generics ( #( #inputs ),* #variadic ) #output {
                    println!("begin");
                    let result = #block;
                    println!("end");
                    return result;
                }
            };
            gen.into()
        }
        _ => panic!("Only fn is allowed!"),
    }
}

嘿,我知道这段代码很丑, but it works 。

属性还可以标记许多的元素,具体可以参见 syn::Item 的定义。 syn::Item 是一个 Enum ,里面的元素除了上面提到的 syn::Item::Fn 之外,还有 mod 定义、 struct/enum/union/trait 定义、 impl 实现甚至 macro_rules! 宏定义等等,只有我想不到的……

除此之外还可以给属性设置参数。还记得我上面代码中无视的 _head: TokenStream 吗,这就是在属性中作为参数出现那一部分,比如在 rocket 库中就用属性来确定标记的函数处理的是哪个路径的路由,如

#[get("/hello/<name>/<age>")]
fn hello(name: String, age: u8) -> String {
    format!("Hello, {} year old named {}!", age, name)
}

于是这个字符串字面量就会被作为第一个参数传入过程宏函数中,并通过 rocket 库的处理来匹配到函数参数中去。因为它也是一个 TokenStream ,因此处理方式是一样的,就不多做介绍了。

参考资料

发表评论