Search code examples
parsingrustmacrosrust-proc-macros

Transforming attributes into identifiers on proc macro derive


I'm writing my first proc macro and despite trying to read through the source for thiserror, structopt and derive_more I can't seem to find exactly what I'm looking for. I want to transform this:

#[derive(Attach)]
#[attach(foo(SomeType, OtherType))]
#[attach(bar(OtherType))]
struct Plugin {}

into this:

impl Attach for Plugin {
    fn attach(self, srv: &mut Server) {
        let this = Arc::new(self);
        srv.foo_order(this.clone(), &[TypeId::of::<SomeType>(), TypeId::of::<OtherType>()]);
        srv.bar_order(this, &[TypeId::of::<OtherType>()]);
    }
}

I've started writing a proc macro but got bungled up trying to parse the attributes:

extern crate proc_macro;

use proc_macro::{Span, TokenStream};
use quote::quote;
use std::any::TypeId;
use syn::{
    parse::ParseStream, parse_macro_input, Attribute, AttributeArgs, DeriveInput, Ident, Result,
};

#[proc_macro_derive(Attach, attributes(attach))]
pub fn register_macro(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    impl_register(input)
}

fn impl_register(input: DeriveInput) -> TokenStream {
    let name = &input.ident;
    let attrs = &input.attrs;

    for attr in attrs {
        if attr.path.is_ident("attach") {
            parse_attach_attribute(&attr);
        }
    }
    println!("{:#?}", input);

    TokenStream::from(quote! {
        impl ::corten::Attach for #name {
            fn attach(self, srv: &mut ::corten::Server) {

            }
        }
    })
}

fn parse_attach_attribute(attr: &Attribute) -> Result<Dependency> {
    let list: syn::MetaList = attr.parse_args()?;
    // println!("{:#?}", list);
    let ident = list.path.get_ident().expect("expected identifier");
    let method = Ident::new(&format!("{}_order", ident), ident.span());
    println!("{:#?}", method);
    let dependencies = list
        .nested
        .into_pairs()
        .map(|pair| pair.into_value())
        .collect::<Vec<_>>();
    // How does one get the identifiers out of a NestedMeta?
    println!("{:#?}", dependencies);
    // attr.parse_args_with(|input: ParseStream| {
    //     let ident = input.p
    //     let method = Ident::new(&format!("{}_order", ident), Span::call_site());
    //     let dep = Dependency {
    //         method,

    //     }
    // })
    unimplemented!()
}
struct Dependency {
    method: Ident,
    dependencies: Vec<Ident>,
}

the difficulty I'm having is how to actually get the attributes list into a usable form? As far as I can tell, I need to get the "foo" and "bar" parsed from &[Attribute] so I can construct the method identifier, and also the "SomeType" & "OtherType" identifiers, which I'll all eventually feed into quote!. If I print the TokenStream in the console, all of the information is there:

[
    Attribute {
        pound_token: Pound,
        style: Outer,
        bracket_token: Bracket,
        path: Path {
            leading_colon: None,
            segments: [
                PathSegment {
                    ident: Ident {
                        ident: "attach",
                        span: #0 bytes(103..109),
                    },
                    arguments: None,
                },
            ],
        },
        tokens: TokenStream [
            Group {
                delimiter: Parenthesis,
                stream: TokenStream [
                    Ident {
                        ident: "predecode",
                        span: #0 bytes(110..119),
                    },
                    Group {
                        delimiter: Parenthesis,
                        stream: TokenStream [
                            Ident {
                                ident: "SomeType",
                                span: #0 bytes(120..128),
                            },
                            Punct {
                                ch: ',',
                                spacing: Alone,
                                span: #0 bytes(128..129),
                            },
                            Ident {
                                ident: "OtherType",
                                span: #0 bytes(130..139),
                            },
                        ],
                        span: #0 bytes(119..140),
                    },
                ],
                span: #0 bytes(109..141),
            },
        ],
    },
    Attribute {
        pound_token: Pound,
        style: Outer,
        bracket_token: Bracket,
        path: Path {
            leading_colon: None,
            segments: [
                PathSegment {
                    ident: Ident {
                        ident: "attach",
                        span: #0 bytes(145..151),
                    },
                    arguments: None,
                },
            ],
        },
        tokens: TokenStream [
            Group {
                delimiter: Parenthesis,
                stream: TokenStream [
                    Ident {
                        ident: "preresolve",
                        span: #0 bytes(152..162),
                    },
                    Group {
                        delimiter: Parenthesis,
                        stream: TokenStream [
                            Ident {
                                ident: "OtherType",
                                span: #0 bytes(163..172),
                            },
                        ],
                        span: #0 bytes(162..173),
                    },
                ],
                span: #0 bytes(151..174),
            },
        ],
    },
]

But I don't have a way to actually get at it. How do I get to tokens[0].stream.ident?


Solution

  • After much messing around I think I have something that works, although I'm happy to accept other answers that are better as I feel this is a bit messy:

    extern crate proc_macro;
    
    use proc_macro2::{Span, TokenStream};
    use quote::quote;
    use syn::{parse_macro_input, Attribute, DeriveInput, Ident, Result};
    
    #[proc_macro_derive(Attach, attributes(attach))]
    pub fn register_macro(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
        let input = parse_macro_input!(input as DeriveInput);
    
        proc_macro::TokenStream::from(impl_register(input))
    }
    
    fn impl_register(input: DeriveInput) -> TokenStream {
        let name = &input.ident;
        let attrs = &input.attrs;
        // println!("{:#?}", input);
    
        let attrs = attrs
            .iter()
            .filter(|attr| attr.path.is_ident("attach"))
            .map(|attr| parse_attach_attribute(&attr).expect("parse failed"))
            .map(|dep| {
                let method: Ident = dep.method;
                let dependencies = dep.dependencies.iter().map(|ident: &Ident| {
                    quote! {
                        std::any::TypeId::of::<#ident>()
                    }
                });
                quote! {
                    srv.#method::<#name, _>(Arc::clone(&this), &[ #(#dependencies),* ]);
                }
            });
    
        quote! {
            impl corten::Attach for #name {
                fn attach(self, srv: &mut corten::Server) {
                    let this = std::sync::Arc::new(self);
                    #(#attrs)*
                }
            }
        }
    }
    
    fn parse_attach_attribute(attr: &Attribute) -> Result<Dependency> {
        let list: syn::MetaList = attr.parse_args()?;
        // println!("{:#?}", list);
        let ident = list.path.get_ident().expect("expected identifier");
        let method = Ident::new(&format!("{}_order", ident), Span::call_site());
        println!("{:#?}", method);
        let dependencies = list
            .nested
            .into_pairs()
            .map(|pair| pair.into_value())
            .filter_map(|pair| match pair {
                syn::NestedMeta::Meta(meta) => match meta {
                    syn::Meta::Path(path) => path.get_ident().cloned(),
                    _ => panic!("only path meta supported"),
                },
                _ => panic!("lit not supported"),
            })
            .collect();
        println!("{:#?}", dependencies);
    
        Ok(Dependency {
            method,
            dependencies,
        })
    }
    struct Dependency {
        method: Ident,
        dependencies: Vec<Ident>,
    }