Transforming attributes into identifiers on proc macro derive

3.3k views Asked by At

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?

1

There are 1 answers

0
leshow On BEST ANSWER

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>,
}