Syn get flat token stream

Question

For the sake of this example: I want to write a proc macro attribute which modifies all number literals in a function body.

I have a syn::Block from a syn::FnItem. I want to map all over the tokens in the block and change literals. I currently have:

let token_stream = func
    .block
    .to_token_stream()
    .into_iter()
    .map(|token| ...)

func.block = Box::new(syn::parse2::<Block>(token_stream.collect()).unwrap());

The problem is to_token_stream from Quote::ToTokens returns a proc_macro2::TokenStream and its into_iter item is proc_macro2::TokenTree which is a nested representation of the tokens.

So for the block representing 2 + (3 + 4)

The iterator is over the following single item:

Group {
    delimiter: Brace,
    stream: TokenStream [
        Literal {
            kind: Integer,
            symbol: "2",
            suffix: None,
            span: #0 bytes(209..210),
        },
        Punct {
            ch: '+',
            spacing: Alone,
            span: #0 bytes(211..212),
        },
        Group {
            delimiter: Parenthesis,
            stream: TokenStream [
                Literal {
                    kind: Integer,
                    symbol: "3",
                    suffix: None,
                    span: #0 bytes(214..215),
                },
                Punct {
                    ch: '+',
                    spacing: Alone,
                    span: #0 bytes(216..217),
                },
                Literal {
                    kind: Integer,
                    symbol: "4",
                    suffix: None,
                    span: #0 bytes(218..219),
                },
            ],
            span: #0 bytes(213..220),
        },
    ],
    span: #0 bytes(203..222),
}

Where the tokens are deeply nested in the structure.

I need a flat representation that would look like:

Literal("2")
Plus
ParenthesisOpen
Literal("3")
Plus
Literal("4")
ParenthesisClose

Is this possible? I could write my own flat_map thinh to go before the map to make it work but that is a lot of work. Does syn expose anything to do this? Maybe I should not parse it into a syn::Block before I do the transform...?

Answer 1

Turns out that syn has a mutable visiting module. Rather than touching the token stream you can use a recursive visitor like so:

use syn::{visit_mut::{self, VisitMut}, LitInt};

struct MyNumberLiteralModifier;

impl VisitMut for MyNumberLiteralModifier {
    fn visit_lit_int_mut(&mut self, node: &mut LitInt) {
        *node = LitInt::new("10", node.span());   
    }
}

MyNumberLiteralModifier.walk_block(&mut my_block);

To do so "visit-mut" feature must be enabled for syn