    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89

    use crate::{Error, Result};
    use proc_macro2::{Delimiter, Ident, Span, TokenStream, TokenTree};
    use quote::ToTokens;
    use std::iter::Peekable;

    pub(crate) fn parse_input(
        input: TokenStream,
    ) -> Result<(Vec<Attribute>, Vec<TokenTree>, TokenTree)> {
        let mut input = input.into_iter().peekable();
        let mut attrs = Vec::new();

        while let Some(attr) = parse_next_attr(&mut input)? {
            attrs.push(attr);
        }

        let sig = parse_signature(&mut input);
        let body = input.next().ok_or_else(|| {
            Error::new(
                Span::call_site(),
                "`#[proc_macro_error]` can be applied only to functions".to_string(),
            )
        })?;

        Ok((attrs, sig, body))
    }

    fn parse_next_attr(
        input: &mut Peekable<impl Iterator<Item = TokenTree>>,
    ) -> Result<Option<Attribute>> {
        let shebang = match input.peek() {
            Some(TokenTree::Punct(ref punct)) if punct.as_char() == '#' => input.next().unwrap(),
            _ => return Ok(None),
        };

        let group = match input.peek() {
            Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::Bracket => {
                let res = group.clone();
                input.next();
                res
            }
            other => {
                let span = other.map_or(Span::call_site(), |tt| tt.span());
                return Err(Error::new(span, "expected `[`".to_string()));
            }
        };

        let path = match group.stream().into_iter().next() {
            Some(TokenTree::Ident(ident)) => Some(ident),
            _ => None,
        };

        Ok(Some(Attribute {
            shebang,
            group: TokenTree::Group(group),
            path,
        }))
    }

    fn parse_signature(input: &mut Peekable<impl Iterator<Item = TokenTree>>) -> Vec<TokenTree> {
        let mut sig = Vec::new();
        loop {
            match input.peek() {
                Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::Brace => {
                    return sig;
                }
                None => return sig,
                _ => sig.push(input.next().unwrap()),
            }
        }
    }

    pub(crate) struct Attribute {
        pub(crate) shebang: TokenTree,
        pub(crate) group: TokenTree,
        pub(crate) path: Option<Ident>,
    }

    impl Attribute {
        pub(crate) fn path_is_ident(&self, ident: &str) -> bool {
            self.path.as_ref().map_or(false, |p| *p == ident)
        }
    }

    impl ToTokens for Attribute {
        fn to_tokens(&self, ts: &mut TokenStream) {
            self.shebang.to_tokens(ts);
            self.group.to_tokens(ts);
        }
    }