Procedural Macros
I’m writing this post because I just finished writing my first procedural, derive macro but didn’t feel like there was enough material on the subject. I’m hoping that this basic example will help others who are doing this for the first time, as it’s pretty nuanced. Macros are in the “Advanced” section of the Rust book after all.
What we’re making
I’m still working on that videogame I told y’all about a few months back, and I’m still doing it in Rust. Although, for the simple things like the main menu, I’m writing more gdscript because I’m wasting too much time on things I shouldn’t be. Anyway, here it is.
#[derive(TimerAttr)]
struct ExampleTimerTwo {
// This timer attribute designates this field as a timer and will generate a poll-like function
// that tests for readiness. If the timer is ready, return true, if the timer is pending, return false.
// You can't return enums, but more on that later.
#[timer]
shield_timer: f64,
#[timer]
laser_timer: f64,
other_field: String,
}
This is a procedural macro. Of the procedural macros, there are 3 sub-types of macros:
- Custom #[derive] macros that specify code added with the derive attribute used on structs and enums
- Attribute-like macros that define custom attributes usable on any item
- Function-like macros that look like function calls but operate on the tokens specified as their argument
This blog post focuses on the #[derive] macro.
The first thing we need to do is create a new cargo package for this. A procedural macro needs to be in it’s own crate. Since this is part of a larger project for me, I leveraged the Cargo Workspaces. Notice that there’s only one Cargo.lock at the root of my workspace. Cargo workspaces use one target directory and one lock file.
Cargo.lock
Cargo.toml
target
kitoshyk
├── Cargo.toml
├── src
└── tests
timer_derive
├── Cargo.toml
├── src
└── tests
It’s a convention to name the procedural macro that implements a derive macro to suffix it with a _derive.
In timer_derive/Cargo.toml
[package]
# according to the Rust book, the naming convention is thing_derive for attribute like procedural macros.
name = "timer_derive"
version = "0.1.0"
authors = ["Drew Gonzales <drewgonzales360@gmail.com>"]
edition = "2018"
[lib]
# This is an important flag that needs to get turned on. It'll put additional rules in the compiler for proc-macros
proc-macro = true
[dependencies]
quote = "1" # This quote library has a nice macro that we'll use
syn = { version = "1", features = ["full"] }
proc-macro2 = "*"
Now let’s take a look at src/lib.rs.
#[macro_use]
extern crate quote;
extern crate proc_macro;
extern crate syn;
use proc_macro::TokenStream;
#[proc_macro_derive(TimerAttr, attributes(timer))]
pub fn timer_macro(_input: TokenStream) -> TokenStream {
// Don't make any additions to the code.
TokenStream::new()
}
Then let’s throw the following in tests/lib.rs
extern crate timer_derive;
use timer_derive::TimerAttr;
#[derive(TimerAttr)]
struct ExampleTimer {
#[timer]
field: f64,
#[allow(dead_code)]
non_timer_field: String,
}
$ cargo test -p timer_derive and it’ll compile.
There are several steps to compiling a piece of code. The first step, is Tokenization. During this stage, the compiler will determine what an integer, identifier (variable names and function names), and keyword (fn, match) is. The incoming parameter to our macro is a TokenStream. This TokenStream has all the info about what structs, variables, and functions are available to us.
We can’t modify the code that comes in, but we can make additions by returning a non-empty TokenStream.
Reading the TokenStream
First, we need to make sure that this attribute only gets added to a struct. To do this, we’ll need to read what’s in the TokenStream. We need to parse the TokenStream into a syn::Item
#[proc_macro_derive(TimerAttr, attributes(timer))]
pub fn timer_macro(input: TokenStream) -> TokenStream {
// The abstract syntax tree gets parsed after the tokenizer finishes its work.
// Here, we're parsing the tree into an "Item", although I could have used anything
// else.
let item: syn::Item = syn::parse(input).expect("failed to parse input");
TokenStream::new()
}
The docs for item say that it can parse any scope {} of code. In our case, it’ll parse the following.
#[derive(TimerAttr)]
struct ExampleTimer {
#[timer]
field: f64,
#[allow(dead_code)]
non_timer_field: String,
}
It’ll include the other struct level attributes if we need to. Then we can pattern match on the output.
#[proc_macro_derive(TimerAttr, attributes(timer))]
pub fn timer_macro(input: TokenStream) -> TokenStream {
// The abstract syntax tree gets parsed after the tokenizer finishes its work.
// Here, we're parsing the tree into an "Item", although I could have used anything
// else.
let item: syn::Item = syn::parse(input).expect("failed to parse input");
let q = match item {
Item::Struct(ref struct_item) => {
TokenStream::new()
}
_ => {
// This block is going to throw a compiler error for us. I'm not exacltly sure what the span is
// but I understand it to represent where in the code the error would be. So if a Struct is the
// item, the span would be the struct and all it's inner fields.
quote_spanned! {
item.span() =>
compile_error!("timer is only valid on structs");
};
TokenStream::new()
}
};
q.into()
}
Since we need to generate functions for all the fields with the #[timer] attribute, we need to iterate through all the fields and find the ones with the #[timer] attribute. This should feel much more familiar to people, so I won’t paste the entire function again.
// Grab the fields of the struct, then filter out any field that doesn't have the
// #[timer] attribute above it.
let timer_fields: Vec<&Field> = struct_item
.fields
.iter()
.filter(|field| {
field.attrs.iter().any(|attr| {
attr.path.get_ident().map(|i| i.to_string()).unwrap() == "timer"
})
})
.collect();
Now that we know which fields are timers from reading the TokenStream, we need to generate code based on the information we just read.
Generating Code
Recall that the end goal is to generate a few functions so that the final TokenStream the compiler see represents the following. I’ve added a second timer field.
#[derive(TimerAttr)]
struct ExampleTimer {
#[timer]
field_one: f64,
#[timer]
field_two: f64,
#[allow(dead_code)]
non_timer_field: String,
}
impl ExampleTimer {
pub fn field_one_is_ready(&self) -> bool {
self.field_one == 0f64
}
pub fn field_two_is_ready(&self) -> bool {
self.field_two == 0f64
}
pub fn decrement_timers(&mut self, delta: f64 ) {
self.field_one = (self.field_one - delta).max(0f64);
self.field_two = (self.field_two - delta).max(0f64);
}
}
To generate an is_ready() function, we’ll just need the field name
/// For a given field, we need a function that checks that the field is ready.
fn implement_poll(field: &Field) -> TokenStream {
let field_name = field.clone().ident.unwrap();
// I'm not sure what the Span is or how it relates to hygienic macros.
let fn_name = Ident::new(&format!("{}_is_ready", field_name,), Span::call_site());
// Here is where the magic happens. The `#` is used to take information from the code we
// actually wrote. IE #field_name came from the var `let field_name` on the first line
// of the function. This quote! macro will turn everything into a proc_macro2::TokenStream
// and we can just coerce it into a proc_macro::TokenStream.
let q = quote! {
pub fn #fn_name(&self) -> bool {
self.#field_name == 0f64
}
};
q.into()
}
Where it gets a little tricky is the mutliple repeated lines in decrement_timers()
/// This function will decrement all the fields with a #[timer] attribute. It'll evaluate to
/// a token stream representing the following.
fn implement_decrement_timer(fields: Vec<&Field>) -> TokenStream {
let field_names: Vec<Ident> = fields.iter().filter_map(|v| v.ident.clone()).collect();
// The #(...)* is explained in https://docs.rs/quote/1.0.7/quote/macro.quote.html#interpolation
let q = quote! {
pub fn decrement_timers(&mut self, delta: f64 ) {
#(
self.#field_names = (self.#field_names - delta).max(0f64);
)*
}
};
q.into()
}
This kind of interpolation isn’t super intuitive, so I think it’s helpful to get as many examples as possible out there into the world. The two blocks of code only represent the functions we’re creating, but don’t say anything about putting them in an impl block. For that, we’ll write another function.
/// Token streams are composable, so we can take the output of the first two, and throw them into
/// a bigger token stream that organizes the functions better.
fn implement(
struct_name: Ident,
polls: Vec<TokenStream>,
increment_timer: TokenStream,
) -> TokenStream {
// First, we need to turn the proc_macro::TokenStreams into proc_macro2::TokenStreams because
// the proc_macro2 crate allows us to use the quote! macro much easier.
let p: proc_macro2::TokenStream = {
let ts = polls
.iter()
.map::<proc_macro2::TokenStream, _>(|f: &TokenStream| f.clone().into())
.into_iter();
proc_macro2::TokenStream::from_iter(ts)
};
let it: proc_macro2::TokenStream = increment_timer.into();
// Now, throw it into an impl block. quote! plays really nicely with proc_macro2::TokenStream (not proc_macro::TokenStream)
// I recommend always using the proc_macro2 where you can. They can be mapped to and fro easily.
let q = quote! {
impl #struct_name {
#p
#it
}
};
q.into()
}
I’ve put the final file in a Gist. After putting all of that together you should be able to use the functions freely without defining them in code.
extern crate timer_derive;
use timer_derive::TimerAttr;
#[derive(TimerAttr)]
struct ExampleTimerTwo {
#[timer]
shield_timer: f64,
#[timer]
laser_timer: f64,
}
#[test]
fn test_increment_multiple_timers() {
let mut e = ExampleTimerTwo {
shield_timer: 39f64,
laser_timer: 10f64,
};
e.decrement_timers(13f64);
assert_eq!(e.shield_timer, 26f64);
assert_eq!(e.laser_timer, 0f64);
assert!(!e.shield_timer_is_ready());
assert!(e.laser_timer_is_ready());
}
Final Thoughts
I’m generally not a fan of autogenerating code like this. In my last job, we had these massive Python scripts that would generate more Python. This was terrible because we obviously didn’t have a compiler, so we had to wait for integration tests or production to find errors. The Python language server also isn’t nearly as strong as the one for Rust. I couldn’t use PyCharm to click through to definitions and had to ctrl+f like a caveman. I don’t think the RLS is perfect, but it’s pretty damn good.
My recommendation is that these kinds of macros should be as simple as possible since we’re abstracting code away and adding code that isn’t there in the original structs definition.
Other Resources: