Search code examples
rustmacrosmetaprogramming

How can I write an (ideally declarative) Rust macro that traverses an AST?


I am trying to write a macro that takes (and expands to) a block of code, but will also collect up the doc strings it found (using the fact that doc strings are just attributes). So for example, the following code:

doc_grab! {
    /// Struct A
    struct A {
        /// Field x
        x: u32,
    }
}

Might expand to:

const A_DOC: ... = ... "Struct A"
                   ... "Field x" ...;

[doc = "Struct A"]
struct A {
    [doc = "Field x"]
    x: u32,
}

I know that this sort of thing is possible, at least with procedural macros, as the Clap library does this using the derive API. What I want to know, however, is it possible to do this with declarative macros, or is it better to stick to doing this with procedural macros.

As a start, I tried making a simple macro walker!, that does nothing but expand (recursively) to the structure it is matching on. So, I came up with this:

macro_rules! walker {
    () => {};
    ($body:tt $($rest:tt)*) => {
        $body walker!($($rest)*)
    };
}

walker! {
    fn get_three(x): u8 {
        1 + 2
    }
}

Unfortunately, that gave me the following error:

error: expected one of `(` or `<`, found `!`
  --> src/main.rs:60:21
   |
60 |           $body walker!($($rest)*)
   |                       ^ expected one of `(` or `<`
...
64 | / walker! {
65 | |     fn get_three(x): u8 {
66 | |         1 + 2
67 | |     }
68 | | }
   | |_- in this macro invocation
   |
   = note: this error originates in the macro `walker` (in Nightly builds, run with -Z macro-backtrace for more info)

Using rust-analyzer I see that its expanding the first fn, then immediately getting stuck on the rest of the function syntax:

// Recursive expansion of walker! macro
// =====================================

fn walker
!(get_three(x):u8 {
  1+2
})

While I could make an explicit case for fn, as well as every other core form in the language, that seems like it would get quite hairy once other macros are involved.

So, is there any way I can walk the syntax tree with declarative macros? Or am I better off going to procedural macros? And if I do need to use procedural macros, what would be the best way to parse the expanded syntax tree? (Like, in Racket, I would just use something like local-expand, but I can't find anything like it in Rust.)


Solution

  • Yes, you can. But, the macro cannot leave invalid syntax outside the macro at any point, as macros are only allowed at certain points in the AST. (Rust cannot stop whatever it is currently parsing whenever it sees an exclamation mark)

    You can get around this by simply "eating" the token and not re-emitting it, with a definition like this:

    macro_rules! walker {
        () => {};
        ($token:tt $($rest:tt)*) => {
            walker! {
                $($rest)*
            }
        }
    }
    

    This means that

    walker! {
        fn get_three() -> u8 {
            3
        }
    }
    

    expands to

    walker! {
        get_three() -> u8 {
            3
        }
    }
    

    which expands to

    walker! {
        () -> u8 {
            3
        }
    }
    

    and so on until there are no more tokens.

    Of course, that means that the macro will eat up all of the tokens, and because of that, the original function definition will no longer be there.

    You can fix this by re-emitting all of the tokens (in a complete, valid token tree) at the start. To do this, you have to have a public macro (walker) which re-emits the tokens and passes them to another macro, which can eat the duplicated tokens. That might look a little something like this:

    macro_rules! walker {
        ($($rest:tt)*) => {
            $($rest)*
            
            _walker_inner! { $($rest)* }
        };
    }
    
    macro_rules! _walker_inner {
        () => {};
        ($token:tt $($rest:tt)*) => {
            _walker_inner! {
                $($rest)*
            }
        }
    }
    

    With this system,

    walker! {
        fn get_three() -> u8 {
            3
        }
    }
    
    fn main() {
        println!("{}", get_three());
    }
    

    will expand to

    fn get_three() -> u8 {
        3
    }
    
    _walker_inner! {
        fn get_three() -> u8 {
            3
        }
    }
    
    fn main() {
        println!("{}", get_three());
    }
    

    which expands to

    fn get_three() -> u8 {
        3
    }
    
    _walker_inner! {
        get_three() -> u8 {
            3
        }
    }
    
    fn main() {
        println!("{}", get_three());
    }
    

    and so on. Playground

    Whether you are better off with procedural macros is a matter of opinion, I think, but it is definitely possible to do what you are suggesting, with the only pain point being concatenating idents together (like your example at the top). Rust doesn't let you take two separate idents and combine them in a declarative macro.

    You will need to do something like re-export the paste crate to ensure that it is usable in dependents.