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!

Answer 1

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.

Answer 2

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.