diff --git a/frb_macros/src/lib.rs b/frb_macros/src/lib.rs index d14c088db8..1cd9232cac 100644 --- a/frb_macros/src/lib.rs +++ b/frb_macros/src/lib.rs @@ -11,10 +11,42 @@ use proc_macro::*; #[proc_macro_attribute] pub fn frb(attribute: TokenStream, item: TokenStream) -> TokenStream { let mut output = format_frb_attribute(format!("#[frb({attribute})]")); + let item = strip_frb_attr(item); output.extend(item); output } +fn strip_frb_attr(item: TokenStream) -> TokenStream { + item.into_iter() + .scan(None, |pound, tok| { + use TokenTree as TT; + match (£, &tok) { + (None, TT::Punct(punct)) if punct.as_char() == '#' => { + *pound = Some(tok); + Some(TokenStream::new()) + } + (Some(_), TT::Group(group)) if is_frb_bracket(group) => { + _ = pound.take(); + Some(format_frb_attribute(format!("#[{}]", group.stream()))) + } + (_, TT::Group(group)) => Some( + [ + pound.take(), + Some(TT::Group(Group::new( + group.delimiter(), + strip_frb_attr(group.stream()), + ))), + ] + .into_iter() + .flatten() + .collect(), + ), + _ => Some(tok.into()), + } + }) + .collect() +} + fn is_frb_bracket(group: &Group) -> bool { matches!((group.delimiter(), group.stream().into_iter().next()), (Delimiter::Bracket, Some(TokenTree::Ident(ident))) if ident.to_string() == "frb") }