How to modify all string literals with a proc_macro_attribute in Rust?

Question

I'm toying around with Rust's procedural macros, and I wanted to make a custom funny attribute that turns all string literals in it's TokenStream into SHOUTING CASE.

Here's what I have so far in my proc-macro lib called amplify:

// lib.rs
extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, File};

#[proc_macro_attribute]
pub fn amplify(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as File);

    // TODO: modify string literals (LitStr) in token tree to be upper case

    let quoted = quote! {#input};
    println!("{:#?}", quoted);  // just print at compile time to see what is produced
    quoted.into()
}

When applied to a simple hello-world program, it should turn the "Hello, World!" string literal into "HELLO, WORLD!" at compile time. Example:

use amplify::amplify;

#[amplify]
pub fn main() {
    println!("Hello, World!");
}

And then running the code:

cargo run
# prints: HELLO, WORLD!

Checking the strings in the binary:

strings target/debug/hello-world | grep -i hello
# prints: HELLO, WORLD! (and possibly other string garbage from the binary)

It is important to note that I want the attribute macro to recursively walk the token tree looking for all string literals. I don't know how to check whether each Item in the tree is a string literal or not, or whether the Item needs to be recursed into.

Any help would be greatly appreciated!

score 5 · Accepted Answer · answered Mar 12 '22 at 10:08

I do not have a lot of experience using proc_macro, but based on this answer I found it was easy to adapt it to manually replace literals in the token tree without the use of the syn or quote crates. This approach involved using the litrs crate to separate out the types of literals.

use proc_macro::{TokenStream, TokenTree, Literal};
use litrs::StringLit;


fn replace_uppercase(item: TokenStream) -> TokenStream {
    item.into_iter()
        .map(|x| {
            match x {
                TokenTree::Group(group) => {
                    // Pass items in group back through replace_uppercase
                    let new_group = Group::new(group.deliminer(), replace_uppercase(group.stream()));
                    TokenTree::Group(new_group)
                },
                TokenTree::Literal(literal) => {
                    // The `litrs` crate was an easy and straightforward approach, so I used it to determine the type of input literals
                    if let Ok(input_string) = StringLit::try_from(literal) {
                        // Convert string literals with uppercase versions
                        TokenTree::Literal(Literal::string(input_string.value().to_uppercase()))
                    } else {
                        TokenTree::Literal(literal)
                    }
                },
                v => v,
            }
          })
        .collect()
}

#[proc_macro_attribute]
pub fn amplify(_attr: TokenStream, item: TokenStream) -> TokenStream {
    replace_uppercase(item)
}

To be honest, syn is probably the more correct choice however I found it difficult to read the documentation since the interface largely consists of macros.

I couldn't follow `syn`'s documentation either. I'd love to see someone's implementation using `syn` if they can get it working! (I couldn't) — camercu, Mar 13 '22 at 00:07
Not using `syn` here is fine. It's the de-facto crate for parsing Rust code in proc macros, but it's not needed for simple transformations like that that works on atomic tokens. If you'd had to work with more complex Rust syntax (e.g. add a parameter to all methods), it would be pretty much required. With `syn`, by the way, you traverse the AST with the `visit()`, `visit_mut()` and `fold()` methods (requires the appropriate feature flags to be on). — Chayim Friedman, Mar 13 '22 at 01:39

score 2 · Answer 2 · answered Mar 13 '22 at 02:02

Just for demonstration purposes, here's how it'll look like using syn:

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

#[proc_macro_attribute]
pub fn amplify(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let mut input = syn::parse_macro_input!(item as syn::File);

    struct StringsModifier;
    impl syn::visit_mut::VisitMut for StringsModifier {
        // You can also implement `visit_lit_byte_str_mut()` to transform `b"..."` strings
        fn visit_lit_str_mut(&mut self, lit: &mut syn::LitStr) {
            *lit = syn::LitStr::new(&lit.value().to_uppercase(), lit.span());
        }
    }
    syn::visit_mut::visit_file_mut(&mut StringsModifier, &mut input);

    input.into_token_stream().into()
}

But note that it'll not descend into macros (println!("a{}", "b") will not be transformed at all) since syn does not have a way to parse them.

How to modify all string literals with a proc_macro_attribute in Rust?

2 Answers2