use core::fmt::{self, Display};
use std::{cell::RefCell, fmt::Debug, str::FromStr};
use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use syn::parse::{Parse, ParseStream, Parser};
#[derive(Clone)]
pub struct SnifedEntry {
pub snif_path: syn::Path,
arrow: syn::Token![=>],
brace: syn::token::Brace,
pub item: syn::Item,
}
#[derive(Clone)]
pub struct SnifedEntries {
first_group: syn::token::Bracket,
pub entries: Vec<SnifedEntry>,
second_group: syn::token::Bracket,
pub macro_input: TokenStream,
}
impl SnifedEntries {
#[must_use]
pub fn span(&self) -> proc_macro2::Span {
self.entries
.first()
.map_or_else(Span::call_site, SnifedEntry::span)
}
}
impl SnifedEntry {
#[must_use]
pub fn span(&self) -> proc_macro2::Span {
self.snif_path
.segments
.first()
.map_or_else(Span::call_site, |segment| segment.ident.span())
}
}
pub struct CommaSeparated<T>(pub Vec<T>);
impl From<CommaSeparated<Token>> for Vec<String> {
fn from(value: CommaSeparated<Token>) -> Self {
value.0.into_iter().map(|t| t.to_string()).collect()
}
}
pub enum Token {
Ident(syn::Ident),
Literal(syn::LitStr),
}
impl Display for Token {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Token::Ident(ident) => write!(f, "{ident}"),
Token::Literal(literal) => write!(f, "{}", literal.value()),
}
}
}
impl Debug for Token {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Token::Ident(ident) => write!(f, "Ident({ident:?})"),
Token::Literal(literal) => write!(f, "Literal({:?})", literal.value()),
}
}
}
impl PartialEq<&str> for Token {
fn eq(&self, other: &&str) -> bool {
match self {
Token::Ident(ident) => ident == *other,
Token::Literal(literal) => literal.value() == *other,
}
}
}
#[doc(hidden)] pub trait TokenStreamInto<T> {
fn convert_token_stream(self) -> syn::Result<T>;
}
impl<T: syn::parse::Parse> TokenStreamInto<T> for TokenStream {
fn convert_token_stream(self) -> syn::Result<T> {
T::parse.parse2(self)
}
}
pub trait IntoTokenStream {
fn into_token_stream(self) -> TokenStream;
}
impl IntoTokenStream for String {
fn into_token_stream(self) -> TokenStream {
TokenStream::from_str(&self).unwrap_or_else(|e| {
compile_error(&format!("Failed to convert String to TokenStream: {e}"))
})
}
}
impl IntoTokenStream for TokenStream {
fn into_token_stream(self) -> TokenStream {
self
}
}
impl IntoTokenStream for () {
fn into_token_stream(self) -> TokenStream {
TokenStream::new()
}
}
fn compile_error(text: &str) -> TokenStream {
quote::quote! {
::core::compile_error!(#text)
}
}
#[macro_export]
macro_rules! output_str {
($($tokens:tt)*) => {
$crate::ux::push_output(format!($($tokens)*));
};
}
#[macro_export]
macro_rules! output {
($($tokens:tt)*) => {
$crate::ux::push_output($crate::prelude::quote!($($tokens)*));
};
}
thread_local! {
static COLLECTED_OUTPUT: RefCell<TokenStream> = RefCell::new(TokenStream::new());
}
pub fn push_output(output: impl IntoTokenStream) {
COLLECTED_OUTPUT.with(|collected_output| {
collected_output
.borrow_mut()
.extend(output.into_token_stream());
});
}
#[doc(hidden)]
#[must_use]
pub(crate) fn flush_output(last_part: TokenStream) -> TokenStream {
COLLECTED_OUTPUT.with(|collected_output| {
let mut collected_output = std::mem::take(&mut *collected_output.borrow_mut());
collected_output.extend(last_part);
collected_output
})
}
impl Parse for Token {
fn parse(input: ParseStream) -> syn::Result<Self> {
if input.peek(syn::Ident) {
Ok(Token::Ident(input.parse()?))
} else if input.peek(syn::LitStr) {
Ok(Token::Literal(input.parse()?))
} else {
Err(syn::Error::new(input.span(), "Expected ident or literal"))
}
}
}
impl<T: Parse> Parse for CommaSeparated<T> {
fn parse(input: ParseStream) -> syn::Result<Self> {
let parser = syn::punctuated::Punctuated::<T, syn::Token![,]>::parse_terminated;
let components = parser(input)?;
Ok(CommaSeparated(components.into_iter().collect()))
}
}
impl syn::parse::Parse for SnifedEntry {
fn parse(input: ParseStream) -> syn::Result<Self> {
let path = syn::Path::parse_mod_style(input)?;
let arrow = input.parse()?;
let content;
let brace = syn::braced!(content in input);
let item = content.parse()?;
Ok(SnifedEntry {
snif_path: path,
arrow,
brace,
item,
})
}
}
impl syn::parse::Parse for SnifedEntries {
fn parse(input: ParseStream) -> syn::Result<Self> {
let items_input;
let first_group = syn::bracketed!(items_input in input);
let mut items = Vec::new();
while !items_input.is_empty() {
items.push(SnifedEntry::parse(&items_input)?);
}
let macro_input;
let second_group = syn::bracketed!(macro_input in input);
Ok(SnifedEntries {
first_group,
entries: items,
second_group,
macro_input: macro_input.parse()?,
})
}
}
impl ToTokens for SnifedEntry {
fn to_tokens(&self, tokens: &mut TokenStream) {
self.snif_path.to_tokens(tokens);
self.arrow.to_tokens(tokens);
self.brace.surround(tokens, |tokens| {
self.item.to_tokens(tokens);
});
}
}
impl ToTokens for SnifedEntries {
fn to_tokens(&self, tokens: &mut TokenStream) {
self.first_group.surround(tokens, |tokens| {
for item in &self.entries {
item.to_tokens(tokens);
}
});
self.second_group.surround(tokens, |tokens| {
self.macro_input.to_tokens(tokens);
});
}
}
#[cfg(test)]
mod tests {
use std::str::FromStr;
use super::*;
#[test]
fn test_parse_string() {
let tokens = TokenStream::from_str(" \"123\" ").unwrap();
let into: Token = tokens.convert_token_stream().unwrap();
assert_eq!(into.to_string(), "123");
}
#[test]
fn test_parse_vec() {
let tokens = TokenStream::from_str(" \"1\", \"2\", \"3\" ").unwrap();
let into: CommaSeparated<Token> = tokens.convert_token_stream().unwrap();
assert_eq!(into.0, vec!["1", "2", "3"]);
}
#[test]
fn test_parse_tts() {
let tokens = TokenStream::from_str("123").unwrap();
let into: TokenStream = tokens.clone().convert_token_stream().unwrap();
assert_eq!(into.to_string(), tokens.to_string());
}
#[test]
fn test_parse_syn_type() {
let tokens = TokenStream::from_str("asd").unwrap();
let into: syn::Ident = tokens.convert_token_stream().unwrap();
assert_eq!(into.to_string(), "asd");
}
#[test]
fn test_streaming_output() {
output_str!("foo");
output_str!("bar");
output! {
"baz" };
let output = flush_output(TokenStream::from_str("qux").unwrap());
assert_eq!(output.to_string(), "foo bar \"baz\" qux");
}
}