code-product-lib 0.4.1

use proc_macro2::{token_stream, Delimiter, Group, Ident, Literal, Punct, TokenStream, TokenTree};

use std::{cell::RefCell, rc::Rc};

/// Our intermediate representation of the code to be generated. This can be either
/// `proc_macro2::TokenTree objects` or references to to be expanded later.
#[cfg_attr(feature = "debug", derive(Debug))]
#[derive(Clone)]
pub(crate) enum Token {
    // TokenTree entries
    RustIdent(Ident),
    RustPunct(Punct),
    RustLiteral(Literal),

    // already expanded
    ProcMacroTokenStream(TokenStream),

    // We need our own group types because we need to store Tokens
    ParenthesisGroup(Vec<Token>),
    BraceGroup(Vec<Token>),
    BracketGroup(Vec<Token>),
    NoneGroup(Vec<Token>),

    NamedReference(String),
    IndexedReference(usize),
}

impl Token {
    pub fn from_tokentree(token: TokenTree) -> Self {
        match token {
            TokenTree::Ident(ident) => Token::RustIdent(ident),
            TokenTree::Punct(punct) => Token::RustPunct(punct),
            TokenTree::Literal(literal) => Token::RustLiteral(literal),
            TokenTree::Group(_) => unimplemented!(),
        }
    }

    pub fn from_group(delimiter: Delimiter, tokens: Vec<Token>) -> Self {
        match delimiter {
            Delimiter::Parenthesis => Token::ParenthesisGroup(tokens),
            Delimiter::Brace => Token::BraceGroup(tokens),
            Delimiter::Bracket => Token::BracketGroup(tokens),
            Delimiter::None => Token::NoneGroup(tokens),
        }
    }
}

// make sure that a Token is as small as necessary
#[test]
fn check_size() {
    #[allow(warnings)] // silence warning: variants `RustToken` and `Unused` are never constructed
    enum ReferenceSize {
        RustToken(TokenTree),
        TokenStream(TokenStream), // TokenStream unfortunately bloats a bit here
        Unused,
    }
    assert_eq!(
        std::mem::size_of::<Token>(),
        std::mem::size_of::<ReferenceSize>()
    );
}

#[derive(Clone)]
#[cfg_attr(feature = "debug", derive(Debug))]
pub(crate) struct TokenIter(Rc<RefCell<std::iter::Peekable<token_stream::IntoIter>>>);

// comment
impl TokenIter {
    /// Create a new TokenIter from a TokenStream
    pub fn new(input: TokenStream) -> Self {
        Self(Rc::new(RefCell::new(input.into_iter().peekable())))
    }

    /// Peek at the next token
    pub fn peek(&self) -> Option<TokenTree> {
        self.0.as_ref().borrow_mut().peek().cloned()
    }

    /// Get the next token
    pub fn next(&self) -> Option<TokenTree> {
        self.0.as_ref().borrow_mut().next()
    }

    /// Check if the next token is a punct and matches the given char
    /// If so, consume it and return true, otherwise return false.
    pub fn next_punct_with(&self, c: char) -> bool {
        let mut iter = self.0.as_ref().borrow_mut();
        if matches!(iter.peek(), Some(TokenTree::Punct(p)) if p.as_char() == c) {
            iter.next();
            true
        } else {
            false
        }
    }

    /// Check if the next token is an ident, if so consume it and return it.
    pub fn next_ident(&self) -> Option<Ident> {
        let mut iter = self.0.as_ref().borrow_mut();
        match iter.peek() {
            Some(TokenTree::Ident(ident)) => {
                let ident = ident.clone();
                iter.next();
                Some(ident)
            }
            _ => None,
        }
    }

    /// Check if the next token is a literal, if so consume it and return it.
    pub fn next_literal(&self) -> Option<Literal> {
        let mut iter = self.0.as_ref().borrow_mut();
        match iter.peek() {
            Some(TokenTree::Literal(literal)) => {
                let literal = literal.clone();
                iter.next();
                Some(literal)
            }
            _ => None,
        }
    }

    /// Check if the next token is a group with the given delimiter, if so return true.
    pub fn peek_group_with(&self, d: Delimiter) -> bool {
        let mut iter = self.0.as_ref().borrow_mut();
        matches!(iter.peek(), Some(TokenTree::Group(group)) if group.delimiter() == d)
    }

    /// Check if the next token is a group with the given delimiter, if so consume it and return it.
    pub fn next_group_with(&self, d: Delimiter) -> Option<Group> {
        let mut iter = self.0.as_ref().borrow_mut();
        match iter.peek() {
            Some(TokenTree::Group(group)) if group.delimiter() == d => {
                let group = group.clone();
                iter.next();
                Some(group)
            }
            _ => None,
        }
    }

    /// Check if the next token is a group with any delimiter, if so consume it and return it.
    pub fn next_group(&self) -> Option<Group> {
        let mut iter = self.0.as_ref().borrow_mut();
        match iter.peek() {
            Some(TokenTree::Group(group)) => {
                let group = group.clone();
                iter.next();
                Some(group)
            }
            _ => None,
        }
    }

    /// Check if there is a token available
    pub fn token_available(&self) -> bool {
        self.peek().is_some()
    }

    /// Get the next token, panic if there is none
    pub fn next_token(&self) -> TokenTree {
        self.next().unwrap()
    }

    /// Creates a `TokenStream` from the remaining tokens
    pub fn stream(self) -> TokenStream {
        TokenStream::from_iter(self.0.borrow_mut().clone())
    }
}

pub(crate) trait ExpandToken {
    fn expand_token(&mut self, token: &Token, defs: &[TokenStream]);
}

impl ExpandToken for TokenStream {
    fn expand_token(&mut self, token: &Token, defs: &[TokenStream]) {
        match token {
            Token::RustIdent(ident) => {
                self.extend(std::iter::once(TokenTree::Ident(ident.clone())))
            }
            Token::RustPunct(punct) => {
                self.extend(std::iter::once(TokenTree::Punct(punct.clone())))
            }
            Token::RustLiteral(literal) => {
                self.extend(std::iter::once(TokenTree::Literal(literal.clone())))
            }
            Token::ProcMacroTokenStream(stream) => {
                self.extend(stream.clone());
            }
            Token::ParenthesisGroup(group) => {
                let mut stream = TokenStream::new();
                group.iter().for_each(|entry| {
                    stream.expand_token(entry, defs);
                });
                self.extend(std::iter::once(TokenTree::Group(Group::new(
                    Delimiter::Parenthesis,
                    stream,
                ))));
            }
            Token::BraceGroup(group) => {
                let mut stream = TokenStream::new();
                group.iter().for_each(|entry| {
                    stream.expand_token(entry, defs);
                });
                self.extend(std::iter::once(TokenTree::Group(Group::new(
                    Delimiter::Brace,
                    stream,
                ))));
            }
            Token::BracketGroup(group) => {
                let mut stream = TokenStream::new();
                group.iter().for_each(|entry| {
                    stream.expand_token(entry, defs);
                });
                self.extend(std::iter::once(TokenTree::Group(Group::new(
                    Delimiter::Bracket,
                    stream,
                ))));
            }
            Token::NoneGroup(group) => {
                let mut stream = TokenStream::new();
                group.iter().for_each(|entry| {
                    stream.expand_token(entry, defs);
                });
                self.extend(std::iter::once(TokenTree::Group(Group::new(
                    Delimiter::None,
                    stream,
                ))));
            }
            Token::NamedReference(name) => {
                panic!("undefined reference '{}'", name)
            }
            Token::IndexedReference(index) => {
                let Some(def) = defs.get(*index) else {
                    panic!("undefined reference '{}'", index)
                };
                self.extend(def.clone());
            }
        }
    }
}

impl ExpandToken for Vec<Token> {
    fn expand_token(&mut self, token: &Token, defs: &[TokenStream]) {
        match token {
            Token::RustIdent(ident) => self.push(Token::RustIdent(ident.clone())),
            Token::RustPunct(punct) => self.push(Token::RustPunct(punct.clone())),
            Token::RustLiteral(literal) => self.push(Token::RustLiteral(literal.clone())),
            Token::ProcMacroTokenStream(stream) => {
                self.push(Token::ProcMacroTokenStream(stream.clone()))
            }
            Token::ParenthesisGroup(group) => {
                let mut entries = Vec::new();
                group.iter().for_each(|entry| {
                    entries.expand_token(entry, defs);
                });
                self.push(Token::ParenthesisGroup(entries));
            }
            Token::BraceGroup(group) => {
                let mut entries = Vec::new();
                group.iter().for_each(|entry| {
                    entries.expand_token(entry, defs);
                });
                self.push(Token::BraceGroup(entries));
            }
            Token::BracketGroup(group) => {
                let mut entries = Vec::new();
                group.iter().for_each(|entry| {
                    entries.expand_token(entry, defs);
                });
                self.push(Token::BracketGroup(entries));
            }
            Token::NoneGroup(group) => {
                let mut entries = Vec::new();
                group.iter().for_each(|entry| {
                    entries.expand_token(entry, defs);
                });
                self.push(Token::NoneGroup(entries));
            }
            Token::NamedReference(name) => {
                panic!("undefined reference '{}'", name)
            }
            Token::IndexedReference(index) => {
                let Some(def) = defs.get(*index) else {
                    panic!("undefined reference '{}'", index)
                };
                self.extend(std::iter::once(Token::ProcMacroTokenStream(def.clone())));
            }
        }
    }
}