How to Use Rust Procedural Macros to Replace Panic with syn’s Fold

Procedural macros are a powerful tool for manipulating Rust code. Programmers that want to write these macros will undoubtedly turn to libraries like syn and quote for parsing and outputting token streams. But in more advanced use cases, the standard tools that syn provides might prove lacking, resulting in fragile, repetitive code.

In this article, we will demonstrate these shortcomings through a toy example, where we replace every panic in a function by an Err. We will first show the kind of code you would normally write. Afterwards, we will turn to the Fold trait, which makes the code for this kind of manipulation a lot more elegant.

Example use case: replacing panics

syn, the standard Rust library for parsing procedural macro input, has a wide range of functions that can aid us in generating and transforming code. You can also use its standard functions for handling multiple, recursive manipulations. That will work just fine.

But we have to cover every possible case - and there could be a lot - ourselves. So you might regret your choice of tooling when you see how much code you have to write.

The wiser course of action is to turn to the library’s Fold trait. Hidden behind a feature flag, it is an excellent choice for recursively changing code, offering multiple methods that allow you to ‘hook in’ to specific parts of your input.

To illustrate this trait’s usefulness, we can turn to an example from my book, Write Powerful Rust Macros: using a procedural macro to replace panics with the Err enum variant in functions. The following code snippet shows the panic_to_result macro in action:

#[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());
}

Converting panics into errors is per se an interesting transformation. But the code’s primary goal was to serve as a guiding example for learning proper error handling in procedural macros, providing answers to questions like "How can I return errors from my macro?" "How do I point that error to a specific line in the user’s code?" etc.

We also briefly introduced the various error-handling styles of modern programming languages, including Rust — which relies heavily on Algebraic Data Types. Hence, getting the macro to work for every possible input was never the goal. In fact, the book’s code looks specifically for panics inside if statements. This is convenient because that is where the panic occurs in our example. You can see it in the following snippet from the macro implementation:

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)
}

This implementation has the advantage of going through the structure of a parsed function step by step. For learning purposes, that’s a good thing. The downside is that this solution lacks extensibility. Do we really want to check every possible type of input? Do we want to add a new match for while expressions, and loop, and ... you get the point.

The Fold trait

The syn library has an advanced tool for use cases like this one, where we want to go through our input recursively: the Fold trait.

Fold, which is hidden behind a feature flag, offers (according to the documentation):

"Syntax tree traversal to transform the nodes of an owned syntax tree. Each method of the Fold trait [...] can be overridden to customize the behavior when transforming the corresponding type of node."

That sounds useful. Another trait hidden behind a feature flag, Visit, is similar. But it uses a borrow of the tree and does not return anything — less suitable for what we want to do here.

Fold helps us visit every node in a program’s AST. And because it takes ownership of the syntax tree, we can modify those nodes, eventually resulting in a tree that has been altered to our liking. In our case, what we would like to do is walk the AST tree, replacing every panic with Err. This, as you will see, requires surprisingly little code when using Fold.

Now, start by creating a new Rust library with cargo init --lib and turn it into a procedural macro by setting proc-macro=true in your Cargo.toml. You will also need to add some dependencies

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

[lib]
proc-macro = true

syn and quote are two procedural macro classics, the former helping us with parsing, the latter turning our data into a TokenStream that Rust can work with. The chosen syn features are important because Fold is hidden behind the fold flag, while the fold::fold_stmt function, that we will also be using, is activated by the full’ feature.

Next, we define our entry point, the function panic_to_result, declared as an attribute macro. An attribute macro replaces the existing code with the code (as a stream of tokens) that we return from it as a stream of tokens. Hence, the output we generate here will replace the annotated function’s entire definition.

panic_to_result first transforms the input into an ItemFn, meaning we are expecting a function,. then it uses a custom struct and fold_item_fn to fold the input, returning the result as a TokenStream and passing the result to 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
}

We are now getting closer to the key section of our implementation. But first, let us look at a helper function called extract_panic_content. This function checks whether a given Macro in the annotated function is a panic, by comparing its identifier with "panic" as a string. If it is a panic, we return its content, i.e. everything between the opening and closing parentheses. We will need that content soon enough to build a message for our Err. In every other case, we return None.

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
    }
}

Now, let’s focus on the core of the macro, some 20 lines of code. We implement Fold for our struct and overwrite the fold_stmt method. Because a panic is a Macro when parsed by syn, we want to check that struct using our extract_panic_content helper to verify whether it really is a panic. If it is, we get back the content/message that was previously inside the `panic`, which we quote to create the new Err return value.

Finally, thanks to a little trick with parse2, the generated tokens are converted  into a statement, which we return from the function. A cool thing to notice here is that no type specification is needed here because Rust infers it based on the function’s output type. When we do not have a macro or a panic, we give back the existing Stmt. Finally, fold::fold_stmt is called to make sure syn keeps on folding.

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)
   }
}

Things to note: while we called fold_item_fn in our macro entry point, we did not implement it. The trait has a perfectly fine default implementation. So why did we implement fold_stmt? Because we know syn parses the content of a function into a Block, a vector of statements. That means the entire body consists of statements, including our panics, and it’s there we should go looking for them!

This may raise further questions. Maybe you are now wondering why we did not just implement a fold_macro if that exists. After all, a panic is parsed as a Macro in syn. And that, in fact, is what I originally wanted to do! Only to realize that I would be manipulating a macro and replacing it with an Err, which is definitely not a macro. And, sadly, the fold_macro definition specifies we have to return a macro.

Complete example

Let’s see our code in action. I have altered the previous example to include a loop. In our main, we will test all three possible paths.

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:?}");;
}

As the output shows, what should have thrown wild panics has turned into docile Results. Other things like, for example, nesting if statements work as well.

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

This is not a full-fledged solution for improving error handling with macros. For example, one drawback is that it only works with functions that already return a Result. But it does the job as an example of how the Fold trait and a bit of custom code can accomplish some very powerful things.

Summary

In this article, we have reviewed how you can write advanced macros to step through Rust code and modify it, thanks to Fold.

We have seen how the standard tooling available in the syn crate allows us to easily transform functions. It can, for example, change an occurrence of a panic into an Err. What it lacks is an elegant way to recursively step through the entire function, automatically executing a change in every applicable location.

The Fold and Visit traits, hidden behind feature flags, solve this limitation. Fold can manipulate the abstract syntax tree of an existing function, and is ideal for our use case. It has many methods — with useful, albeit basic, default implementations — that allow you to handle every occurrence of a given type. For example, the fold_macro lets you manipulate every macro within your function. Moreover, the fold_stmt method helped us step through the entire function’s content, changing every panic with minimal effort.