Rust 过程宏：用 syn Fold 优雅替换 Panic

rust| 2024-03-28

图0: rustRust 过程宏：用 syn Fold 优雅替换 Panic

Procedural macros 是操作 Rust 代码的强大工具。编写这些宏的程序员通常会使用像 syn 和 quote 这样的库来解析和输出标记流。然而，在更复杂的用例中，syn 库提供的标准工具可能无法满足所有需求。有时，标准工具的功能显得捉襟见肘，导致代码变得脆弱且充满重复。

本文将通过一个玩具示例来揭示这些不足之处，即我们将替换函数中的每个 panic 为 Err。首先，我们将展示通常的代码写法。然后，我们将引入 Fold trait，展示它如何使这种操作的代码变得更加优雅。

示例用例：替换 panic

syn 作为 Rust 中解析过程宏输入的标准库，其功能丰富，能够助力开发者高效地生成和转换代码。其标准功能在处理单一或简单的递归操作时表现良好，但在面对复杂多样的场景时，开发者往往需要自行处理各种可能出现的情况，这无疑增加了工作量和出错的可能性。

当我们面对需要编写大量重复代码以处理不同情况时，可能会开始质疑选择 syn 是否明智。为了克服这一挑战，我们可以转向 syn 库中的 Fold trait。尽管这个 trait 被隐藏在某个特性标志之下，但它在递归地改变代码结构时表现出了强大的能力。Fold trait 提供了多种方法，允许开发者在输入的特定部分进行 “挂钩” 操作。

为了更直观地展示 Fold trait 的实际应用，我们可以参考《Write Powerful Rust Macros》一书中的例子。在这个例子中，我们展示了如何使用过程宏将函数中的 panic 调用替换为返回 Err 枚举变体的操作。以下是 panic_to_result 宏的一个简化版代码片段，它展示了这一转换过程的具体实现：

#[panic_to_result] // use our macro
fn create_person_with_result(name: String, age: u32) -> Result<Person, String> {
   if age > 30 {
       panic!("I hope I die before I get old"); // <- panic will be replaced by an Err
   }
   Ok(Person {
       name,
       age,
   })
}


fn main() {
   // the assertion shows we got back an Err instead of a panic!
   assert!(create_person_with_result("name".to_string(), 31).is_err());
}

在现代编程领域中，各种编程语言在处理错误时都展现出了独特的风格。其中，Rust 语言以其对代数数据类型的深度依赖而脱颖而出。值得注意的是，我们并非旨在将宏应用于所有可能的输入情况。实际上，本书所提供的代码更侧重于识别和处理那些存在于 if 语句内部的 panic 情况。这种设计选择极具实用性，因为这正是我们示例中 panic 出现的典型场景。通过以下宏实现的代码片段，你可以清晰地看到这一点：

match expression {
   // check the if expressions for panics
   Expr::If(mut ex_if) => {
       let new_statements: Vec<Stmt> = // modify existing statements
       ex_if.then_branch.stmts = new_statements; // and put them inside the if
       Stmt::Expr(Expr::If(ex_if), token)
   },
   // return all other expressions without modification
   _ => Stmt::Expr(expression, token)
}

Fold trait

syn 库提供了一项强大的工具，即 Fold trait，它尤其适用于我们需要递归遍历输入语法树（AST）的场景。Fold 隐藏在特性标志之后，但根据官方文档描述：

“Fold trait 用于遍历并转换拥有权的语法树节点。其每个方法都可以被重写，以定制在转换相应类型节点时的行为。”这一描述凸显了其潜在的有用性。尽管 syn 库中还有另一个特性标志隐藏的 trait，即 Visit，它与 Fold 类似，但使用树的引用（borrow）并不返回任何结果，因此并不适合我们当前的需求。

Fold 允许我们访问程序 AST 中的每个节点。由于它拥有对语法树的所有权，我们可以对这些节点进行修改，并最终得到一个按我们意愿改造后的树。在我们的案例中，我们的目标是遍历 AST 树，将每个 panic 调用替换为 Err 表达式。你将发现，使用 Fold trait 所需的代码量竟异常之少。

接下来，让我们通过运行 cargo init --lib 命令来创建一个新的 Rust 库，并将其转换为过程宏。这可以通过在 Cargo.toml 文件中设置 proc-macro = true 来实现。此外，我们还需要添加一些必要的依赖项，以支持我们的宏实现。

[dependencies]
quote = "1.0.33"
syn = { version = "2.0.39", features = ["fold", "full"]}


[lib]
proc-macro = true

接下来，我们定义入口点函数 panic_to_result，它是一个属性宏。属性宏的作用在于将其返回的代码（以令牌流的形式）直接替换原有的代码。因此，我们在此生成的输出将完全取代被标记函数的定义。

panic_to_result 首先会将输入转换为一个 ItemFn 类型，这表示我们期望的输入是一个函数定义。随后，它利用一个自定义的结构和 fold_item_fn 方法来折叠输入，并将结果以 TokenStream 的形式返回。最后，我们将这个 TokenStream 传递给 quote 宏，以便生成最终的替换代码。

use proc_macro::TokenStream;
use quote::quote;


#[proc_macro_attribute]
pub fn panic_to_result(_attr: TokenStream, input: TokenStream) -> TokenStream {
   let item: ItemFn = syn::parse(input).unwrap(); // parse the input
   let result = ChangePanicIntoResult.fold_item_fn(item); // fold it
   quote!(#result).into() // and return the result
}

fn extract_panic_content(mac: &Macro) -> Option<TokenStream2> {
    let does_panic = mac.path.segments.iter()
        .any(|v| v.ident.to_string().eq("panic"));


    if does_panic {
        Some(mac.tokens.clone())
    } else {
        None
    }
}

最后，利用 parse2 的巧妙之处，生成的令牌被顺利转换为一个语句，并由函数返回。在此过程中，值得注意的是，这里并不需要显式指定类型规范，因为 Rust 编译器会根据函数的输出类型进行自动推断。当不存在宏或 panic 调用时，我们则直接返回现有的 Stmt 对象。最后，通过调用 fold::fold_stmt，我们确保了 syn 库能够继续对语句进行折叠处理，从而完成整个转换过程。

use quote::quote;
use syn::{fold, ItemFn, Macro, Stmt};
use syn::fold::Fold;


struct ChangePanicIntoResult; // the struct that we were calling in the entry point


impl Fold for ChangePanicIntoResult {
   fn fold_stmt(&mut self, stmt: Stmt) -> Stmt {
       let new_statement: Stmt = match stmt {
           Stmt::Macro(ref mac) => {
               let output = extract_panic_content(&mac.mac); // helper to get the panic message
               output
                   .map(|t| quote! {
                       return Err(#t.to_string());
                   })
                   .map(syn::parse2)
                   .map(Result::unwrap)
                   .unwrap_or(stmt)
           }
           // panics should be inside a 'Macro', so in every other case we return
           _ => stmt
       };
       // keep folding
       fold::fold_stmt(self, new_statement)
   }
}

这确实可能引发一系列更深层次的问题。或许你此刻正疑惑，为何我们没有实现一个 fold_macro 功能（如果它存在的话）。毕竟，在 syn 库中，panic 被解析为一个 Macro。事实上，这曾是我最初的设想！然而，随着对问题的深入理解，我意识到这样的操作实际上并不可行。原因是，如果我们尝试对一个宏进行操作，并将其替换为一个 Err 表达式，那么这样的替换结果本身就不再是一个宏了。更遗憾的是，fold_macro 的定义明确要求我们必须返回一个宏，这使得我们的设想无法实现。

完整示例

让我们深入探究一下我们的代码在实际运行时的效果。我特地对之前的示例进行了调整，加入了循环结构。在我们的主函数中，我们将对三种可能的路径进行详尽的测试。

use fold_macro::panic_to_result;


#[derive(Debug)]
pub struct Person {
   name: String,
   age: u32,
}


#[panic_to_result]
fn create_person_with_result(name: String, age: u32) -> Result<Person, String> {
   // 'if' works
   if age > 30 && age < 50 {
       panic!("I hope I die before I get old");
   }
   // but now loop does as well
   loop {
       if age > 50 {
           panic!("This person is old... very old");
       }
       break
   }
   Ok(Person {
       name,
       age,
   })
}


fn main() {
   let first = create_person_with_result("name".to_string(), 20);
   println!("{first:?}");
   let second = create_person_with_result("name".to_string(), 40);
   println!("{second:?}");
   let third = create_person_with_result("name".to_string(), 51);
   println!("{third:?}");;
}

Ok(Person { name: "name", age: 20 }) Err("I hope I die before I get old") Err("This person is old... very old")

尽管这并非一个全面完善的错误处理宏解决方案，但它确实为解决特定问题提供了一个颇具启发性的示例。该宏的局限性在于，它目前仅适用于那些已经设计为返回 Result 类型的函数。然而，这并不影响它作为一个展示 Fold trait 和自定义代码如何结合实现强大功能的出色案例。

总结

在本文中，我们深入探讨了如何借助 Fold trait 编写高级宏来遍历并修改 Rust 代码。syn crate 提供的标准工具集使得我们能够以简洁高效的方式转换函数。举例来说，我们成功地利用这些工具将 panic 调用替换为 Err 表达式。然而，此前缺乏一种优雅且自动化的方法来递归遍历整个函数，并在每个适用的位置执行更改。

Fold 和 Visit trait 的出现，打破了这一局限。尽管它们隐藏在特性标志之后，但为我们提供了强大的工具。Fold trait 尤其适用于操作函数的抽象语法树（AST），因此非常符合我们的用例。它提供了多种方法，这些方法尽管带有基本的默认实现，但却极具实用性，能够处理给定类型的每个出现。比如，fold_macro 方法允许我们操纵函数中的每个宏。此外，fold_stmt 方法帮助我们以最小的努力遍历整个函数的内容，从而轻松地更改每个 panic。

原文链接：

https://www.infoq.com/articles/rust-procedural-macros-replace-panic/

本文文字及图片出自 InfoQ