Skip to content
On this page

Rust 中的宏编程和实战

Rust 以高性能和安全著称,但是它的宏系统也是非常强大的。宏是一种元编程的手段,可以在编译期间生成代码。

Rust 的宏分为两种:声明式宏和过程宏。声明式宏是 Rust 语言的一部分,而过程宏是一个独立的库。本文将主要以实战过程宏为主了解下过程后工作的流程,更多概念相关的介绍相关的可以参考 https://course.rs/advance/macro.html

声明式宏

声明式宏是 Rust 语言的一部分,它使用 macro_rules! 关键字定义。声明式宏的语法类似于函数,但是它可以在编译期间生成代码。

定义宏

下面是一个简单的宏定义:

rust
macro_rules! say_hello {
  () => {
    println!("Hello, world!");
  };
}

这个宏没有参数,当调用 say_hello! 时,会打印 Hello, world!。Rust 中是不存在重载的,但是宏可以实现类似重载的功能:

rust
macro_rules! say_hello {
  () => {
    println!("Hello, world!");
  };
  ($name:expr) => {
    println!("Hello, {}!", $name);
  };
}

这个宏有两个分支,一个是不带参数的,一个是带参数的。当调用 say_hello! 时,会根据参数的个数选择对应的分支。

你还可以在宏的分支中调用另一个分支:

rust
macro_rules! say_hello {
  () => {
    // 调用另一个分支
    say_hello!("world");
  };
  ($name:expr) => {
    println!("Hello, {}!", $name);
  };
}

生成重复代码

那么声明式宏有什么用呢?一个常见的用途是生成重复的代码。但宏生成代码和利用函数生成代码除了编译期间和运行期间的区别外,还有什么区别呢?

我们可能会遇到一个场景,例如下面有两个函数 ab,他们都有类似的逻辑: do something

rust
fn a() {
  // do something
}

fn b() {
  // do something
}

但他们接受的参数不同,我们就可以利用上面提到的利用声明式宏重载的特性生成抽取重复代码

rust
macro_rules! do_something {
  () => {
    // do something
  };
  ($name:expr) => {
    // do something
  };
}

fn a() {
  do_something!();
}

fn b() {
  do_something!("b");
}

还有一种更有意思的情况也可以让我们利用宏生成重复代码,假设 do_something 参数逻辑上都可以利用函数的方式抽取公共逻辑。但是他在参数上使用了一个很复杂的类型 T: A<B<C<D>>>,更有甚者中间可能还有多个范型和生命周期参数

笔者遇到一个更极端的情况是 T 参数中 B 没有被导出,所以无法直接通过函数定义参数 / 返回值类型,这种情况下我们就可以通过宏来直接跳过参数类型 / 返回值的定义,直接生成代码

rust
macro_rules! do_something {
  ($param_a:ident, $param_b:ident) => {
    // do something
  };
}

fn a() {
  let param_a: T = ...;
  let param_b: B = ...;

  do_something!(param_a, param_b);
}

更多关于声明式宏的内容也可以参考官方文档

过程宏

如果说声明式宏是黑科技的话,那么过程宏就是黑科技中的黑科技。过程宏可以让你在代码中做到解析语法树从而动态生成代码。

我们以最常见的代码展示下它的魔力

rust
// 定义了一个简单的结构体
#[derive(Default)]
struct A {
  a: i32,
  b: i32,
}

fn main() {
  // 调用 `default` 方法
  let a = A::default();
}

这段代码中很多人都会感到困惑,#[derive(Default)] 是什么意思?为什么 A 可以调用 default 方法?

这就是过程宏的魔力,#[derive(Default)] 是一个过程宏,它会在编译期间解析 A 的结构体定义,然后生成一个 default 方法。这个方法会初始化结构体的所有字段为默认值。

相当于

diff
struct A {
  a: i32,
  b: i32,
}

// 以下代码是由 #[derive(Default)] 生成的
+ impl Default for A {
+  fn default() -> Self {
+    Self {
+      a: Default::default(),
+      b: Default::default(),
+    }
+  }
+ }

类函数宏

另外一种也可以展示过程宏的黑科技,我们用 非主流 语法来生成 Github API,指明它的参数 / 返回值甚至是测试用例。

rust
github_api! {
  GithubCommitAPI {
    get_commit {
      path "/repos/{}/{}/commits/{}"
      params {
        owner String
        repo String
        sha String
      }
      response GithubCommit
      test {
        params {
          "panghu-huang"
          "octocrate"
          "123456"
        }
        assert assert_eq!(res.sha, "123456")
      }
    }
  }
}

这部分代码会生成类似下面的代码:

rust
struct GithubCommitAPI {
  api_client: reqwest::Client,
}

impl GithubCommitAPI {
  fn get_commit(&self, owner: impl Into<String>, repo: impl Into<String>, sha: impl Into<String>) -> GithubCommit {
    let response = self
      .api_client
      .get(&format!("/repos/{}/{}/commits/{}", owner.into(), repo.into(), sha.into()))
      .send()
      .await
      .unwrap();

    response.json().unwrap()
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[tokio::test]
  async fn test_get_commit() {
    let api = GithubCommitAPI;
    let res = api.get_commit("panghu-huang", "octocrate", "123456").await.unwrap();

    assert_eq!(res.sha, "123456");
  }
}

这就是过程宏中非常有意思的黑科技,它可以发挥你的各种奇思妙想,定义你自己的语法去写代码。

上面的例子能够充分展示过程宏的魔力,但是对于刚开始接触宏的同学来说太复杂了。接下去我们用个更简单的例子实战下过程宏。

过程宏实战

需求

我们用过程宏实现一个 proto! 宏,在 proto! 中使用 protobuf 的语法定义一个结构体和函数,并生成简单的 impl 代码。

rust
proto! {
  service Greeter {
    rpc SayHello (HelloRequest) returns (HelloResponse) {}
  }
}

这个宏会生成下面的代码:

rust
struct Greeter;

impl Greeter {
  fn say_hello(&self, req: HelloRequest) -> HelloResponse {
    // 实现
  }
}

创建新项目

因为过程宏的特殊性,因此 Rust 要求过程宏必须在一个独立的库中。我们创建一个新的库项目:

bash
cargo new --lib proto-macro

Cargo.toml 中声明这个库是一个过程宏库:

toml
[lib]
proc-macro = true

然后在 proto-macroCargo.toml 中添加依赖:

toml
[dependencies]
quote = "1"
syn = "2"

synquote 都是由 dtolnay 开发的非常强大库,它们分别用于解析和生成 Rust 代码。他们具体的作用在下面的 流程梳理 中我们做个简单的介绍。

流程梳理

过程宏生成代码分为两个重要的步骤:解析输入语法生成输出代码

解析输入语法 就是我们明确知道我们期望的输入是什么,并且这个输入能提供给我们什么信息。在我们的例子中,我们期望的输入是一个 service 定义,它包含一个服务名和一个 rpc 定义。

rust
proto! {
  // 服务定义
  service Greeter {
    // rpc 定义
    rpc SayHello (HelloRequest) returns (HelloResponse) {}
  }
}

其中 service 定义包含一个服务名,rpc 定义包含一个函数名、参数和返回值。我们需要解析这些信息,然后可以使用机构体将这些信息保存下来。

rust
/// Proto 定义
struct Proto {
  service: Service,
}

/// Service 定义
struct Service {
  name: String,
  rpcs: Vec<Rpc>,
}

/// Rpc 定义
struct Rpc {
  name: String,
  request: String,
  response: String,
}

proto! {}struct Proto {} 的解析就是这样的一个过程。而这个过程的实现我们依赖于 syn 这个库,他能让我们更简单的实现语法解析。

生成输出代码 就是我们拿到了输入信息(上面的 struct Proto {}),然后根据这个信息生成对应的 Rust 代码。在这个例子中,就是生成以下代码:

rust
struct Greeter;

impl Greeter {
  fn say_hello(&self, req: HelloRequest) -> HelloResponse {
    // 实现
  }
}

这个过程的实现我们依赖于 quote 这个库,他能让我们更简单的生成 Rust 代码。

解析输入语法

参考 syn 文档得知,首先我们需要定义一个过程宏,它接受一个 TokenStream 输入,并返回一个 TokenStream 输出(就是接受 AST 输入,返回 AST 输出)。

rust
use proc_macro::TokenStream;
use syn::parse_macro_input;

#[proc_macro]
pub fn proto(input: TokenStream) -> TokenStream {
  // 解析输入语法
  let proto = parse_macro_input!(input as Proto);

  // 标记为 TODO,后续实现
  todo!()
}

我们希望最终得到一个 Proto 结构体,因此我们这里用 input as Proto 来解析输入语法。但这里会报错

  --> src/lib.rs:14:43
  |
  |   let proto = parse_macro_input!(input as Proto);
  |                                           ^^^^^ the trait `Parse` is not implemented for `Proto`
  |

这是因为 syn 并不知道如何解析 Proto,我们需要实现 Parse trait 来告诉 syn 如何解析 Proto

rust
use syn::{parse, parse::Parse, parse::ParseStream, Result};

/// 解析 Proto
impl Parse for Proto {
  fn parse(input: ParseStream) -> Result<Self> {
    // 解析语法
  }
}

那么上面的 input as Proto 这句在这里就可以简单的理解为 let proto = Proto::parse(input)

在真正开始解析语法前,我们需要先了解一下 syn 这个库的基本用法。syn 是一个 Rust 语法解析库,它可以解析 Rust 代码并生成一个语法树(AST)。并且它支持让我们 忽略空格逐字拆解

解析的关键就是忽略空格逐字拆解。那么在上面给他上面 proto! 的例子中,忽略空格逐字拆解 后大概是这样子的

1. service              // 固定关键字 service
2. Greeter              // 服务名
3. {                    // 花括号
4. rpc                  // 固定关键字 rpc
5. SayHello             // ...
6. (              
7. HelloRequest
8. )
9. returns
10. (
11. HelloResponse
12. )

将这段转换为伪代码,大概是这样子的

rust
fn parse(input: ParseStream) {
  input.parse('service'); // 解析固定关键字 service
  // 在这里拿到了服务名
  let service_name = input.parse<String>(); // 服务名
  input.parse('{'); // 解析花括号
  input.parse('rpc'); // 解析固定关键字 rpc
  // 在这里拿到了函数名
  let rpc_name = input.parse<String>(); // 函数名
  // ...
}

按照上面的拆解后就非常好理解解析的整个流程了。syn 真正的解析代码其实和这个非常类似。

那么在将上面的伪代码转化成真正的解析代码前,有几个在解析过程中真正涉及到的概念我们需要先了解下

  • Ident:标识符,例如 struct Greeter 中的 Greeter 就是一个标识符 (可以简单类比为 String)。
  • Type:类型,例如 fn say_hello(&self, req: crate::types::HelloRequest) 中的 crate::types::HelloRequest 就是一个类型。
  • Keyword固定关键字,例如 structfn 等。上面的伪代码中的 servicerpc 就是关键字。
  • Token:标记,例如 :; 等。

了解这些后,我们就可以将这些概念代码到真正的代码中了

rust
use syn::{parse, parse::Parse, parse::ParseStream, Result, Ident, Token};

impl Parse for Proto {
  fn parse(input: ParseStream) -> Result<Self> {
    // 解析固定关键字 service
    input.parse::<keyword::service>()?; 
    // 这里换成 `Ident`
    let service_name = input.parse<Ident>()?; 
    // Token
    input.parse::<Token(`{`)>()?; 
    // 解析固定关键字 rpc
    input.parse::<keyword::rpc>()?; 
  }
}

这里真正的代码中体现了几个细节

  • 都是使用 input.parse::<T>()? 来解析,其中 T 就是我们想要的类型,可以是 IdentToken 等。
  • 都是以 ? 结尾,这是因为 parse 方法返回的是 Result 类型,如果我们希望得到关键字 rpc,但是实际上输入的是 rpc1,那么解析就会失败返回 Err

上面的代码真正运行时就会发现了,除了 input.parse<Ident>()? 能够正常运行外,这是因为

  • keyword::servicekeyword::rpc 都是我们自定义的关键字,syn 并不知道如何解析这些关键字。
  • {}() 等不算真正意义上的 Token,所以也无法使用 Token 解析。
自定义关键字

那么我们如何解决上面的问题呢?我们可以通过自定义关键字来解决这个问题。我们可以通过 syn 提供的 custom_keyword! 宏来定义自定义关键字。

rust
use syn::custom_keyword;

// 把关键字定义当在 keyword 模块下
mod keyword {
  custom_keyword!(service);
  custom_keyword!(rpc);
  custom_keyword!(returns);
}

这样我们就可以使用 keyword::servicekeyword::rpc 来解析我们自定义的关键字了。

rust
input.parse::<keyword::service>()?; 
input.parse::<keyword::rpc>()?; 
braced!parenthesized!

对于 {}() 这种 Token,我们可以使用 braced!parenthesized! 来解析。

其中 braced! 用于提取 {} 中的内容,parenthesized! 用于提取 () 中的内容。这两个宏用法类似:

rust
use syn::braced;

// content 就是 {} 中的内容
let content;
braced!(content in input);

通过 braced! 我们就能提取到 {} 中的内容了。具体什么意思呢

假设现在的 input 为 { a b c } d f g,那么在执行 braced!(content in input); 后, content 就是 a b cinput 就是除去 {} (包括 {} 中的内容)后剩下的 d f g

parenthesized! 类似。

那么我们就可以利用 braced!parenthesized! 来解析 {}() 中的内容了。

rust
use syn::{braced, parenthesized, parse, parse::Parse, parse::ParseStream, Result, Ident, Token};

impl Parse for Proto {
  fn parse(input: ParseStream) -> Result<Self> {
    // 解析固定关键字 service
    input.parse::<keyword::service>()?; 
    // 解析 service 名
    let service_name = input.parse<Ident>()?; 

    // 提取 {} 中的内容 -> `rpc SayHello (HelloRequest) returns (HelloResponse) {}`
    let content;
    braced!(content in input);

    // 因为 rpc 定义在 service 的 `{}` 中,因此这里我们需要用 content 而不是 input
    content.parse::<keyword::rpc>()?; 

    let rpc_name = content.parse<Ident>()?;

    // rpc 中请求参数的定义 -> `(HelloRequest)`
    let request_content;
    parenthesized!(request_content in content);

    // 解析请求参数的类型
    let request = request_content.parse<Type>()?;

    // 固定关键字 returns
    content.parse::<keyword::returns>()?;

    // rpc 中返回值的定义 -> `(HelloResponse)`
    let response_content;
    parenthesized!(response_content in content);

    // 解析返回值的类型
    let response = response_content.parse<Type>()?;

    // 注意:到这里并未完成解析,还剩下最后的 `{}`

    // 解析最后的 `{}`
    let _content;
    braced!(_content in content);
  }
}

至此,我们已经完成了解析输入语法的过程。但是我们突然想增加一个新的功能,我们想要允许用户在 proto!package package_name 和多个 rpc 定义

rust
proto! {
  package hello;

  service Greeter {
    rpc SayHello (HelloRequest) returns (HelloResponse) {}
    rpc SayHello2 (HelloRequest) returns (HelloResponse) {}
  }
}

并且我们希望 packageservice 定义的顺序是任意的。那么我们该如何解析?

TODO