chunked_quote_impl/

lib.rs

use proc_macro2::{
	Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree,
	token_stream::IntoIter,
};
use quote::{TokenStreamExt, quote as oquote, quote_spanned as oquote_spanned};

/// append to [`TokenStream`] code generating the input tokens.
///
/// **syntax:** `chunk!(stream: ident, ...)`.
///
/// all not interpolated tokens inherit the [`Span::call_site`] span.
///
/// supported syntax:
/// - `#ident`: append the resolved value of `ident` through [`ToTokens`](quote::ToTokens).
/// - `#{expr}`: append the resolved value of `expr` through [`ToTokens`](quote::ToTokens).
/// - `##`: append a `#`.
/// - `#op expr #{tokens}`: append `tokens` based on the evaluation of `op expr`. `op` can be `if`, `for`, `while`, `else`, `match`.
/// - `#do {expr}`: execute `expr` at its definition point in the strucutre.
/// - other tokens gets appended.
///
/// # example
/// ```
/// let fields = &[
///     (Ident::new("a", Span::call_site()), Ident::new("u32", Span::call_site())),
/// 	(Ident::new("b", Span::call_site()), Ident::new("bool", Span::call_site())),
/// 	(Ident::new("c", Span::call_site()), Ident::new("char", Span::call_site())),
/// 	];
/// let mut stream = TokenStream::new();
/// let public = true;
/// chunk!(stream,
/// 	#if public #{ pub }
/// 	struct Example {
///         #for (field, ty) in fields #{ #field: #ty, }
/// 	}
/// 	impl Example {
///         #do { gen_accessors(stream, fields) }
/// 	}
/// );
/// fn gen_accessors(mut stream: &mut TokenStream, fields: &[(Ident, Ident)]) {
///     chunk!(stream, #for (field, ty) in fields #{
///         fn #{format_ident!("get_{field}")} (&self) -> #ty {
///             self.#field
/// 		}
/// 	});
/// }
/// ```
#[proc_macro]
pub fn chunk(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	chunk_impl(TokenStream::from(input), true, false).into()
}

/// [`chunk`] but with specified span for all tokens not interpolated.
///
/// **syntax:** `chunk_spanned!(stream: ident, span: expr, ...)`
#[proc_macro]
pub fn chunk_spanned(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	chunk_impl(TokenStream::from(input), true, true).into()
}

/// create a [`TokenStream`] from the input tokens.
#[proc_macro]
pub fn quote(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	chunk_impl(TokenStream::from(input), false, false).into()
}

/// [`quote!`] but with specified span for all tokens not interpolated.
///
/// **syntax:** `chunk_spanned!(span: expr, ...)`
#[proc_macro]
pub fn quote_spanned(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	chunk_impl(TokenStream::from(input), false, true).into()
}

/// simplify parsing.
struct Cursor {
	iter: IntoIter,
	/// used for peek
	cur_token: Option<TokenTree>,
	/// the result
	res: TokenStream,
	/// the span after the last token, group end delim
	end_span: Span,
}

impl Cursor {
	/// create a new [`Cursor`]
	fn new(iter: TokenStream, end_span: Span) -> Self {
		Self {
			iter: iter.into_iter(),
			res: TokenStream::new(),
			cur_token: None,
			end_span,
		}
	}
	/// get the next [`TokenTree`]
	fn next(&mut self) -> Option<TokenTree> {
		self.cur_token.take().or_else(|| self.iter.next())
	}
	/// get the next [`TokenTree`] without consuming
	fn peek(&mut self) -> Option<&TokenTree> {
		if self.cur_token.is_none() {
			self.cur_token = self.iter.next();
		}
		self.cur_token.as_ref()
	}
	/// test if the next [`TokenTree`] is a keyword
	fn peek_kw(&mut self, kw: &str) -> bool {
		match self.peek() {
			Some(TokenTree::Ident(ident)) => ident == kw,
			_ => false,
		}
	}
	/// push a chunk into the result
	fn add<T>(&mut self, tokens: impl IntoIterator<Item = T>)
	where
		TokenStream: Extend<T>,
	{
		self.res.extend(tokens);
	}
	/// raise an error of `expected token`
	fn expected(&mut self, expected: &str, span: Option<Span>) {
		let span = span.unwrap_or(self.end_span);
		let msg = &format!("expected {expected}");
		self.res.extend(oquote_spanned! { span => ::core::compile_error!(#msg); });
	}
	/// comsume an identifier, raising an error if not
	fn eat_ident(&mut self) -> Option<Ident> {
		match self.next() {
			Some(TokenTree::Ident(ident)) => return Some(ident),
			t => self.expected("identifier", t.map(|t| t.span())),
		}
		None
	}
	/// comsume a punctuation, raising an error if not
	fn eat_punct(&mut self, char: char) -> bool {
		match self.next() {
			Some(TokenTree::Punct(punct)) if punct.as_char() == char => return true,
			t => self.expected(&format!("`{char}`"), t.map(|t| t.span())),
		}
		false
	}
	/// comsume a brace group, raising an error if not
	fn eat_brace(&mut self) -> Option<Group> {
		match self.next() {
			Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Brace => {
				return Some(g);
			}
			t => self.expected("`(`", t.map(|t| t.span())),
		}
		None
	}
	/// continously eat tokens until `pred` returns `true` or end of input, raising an error if empty.
	fn eat_until(
		&mut self, expected: &str, pred: impl Fn(&TokenTree) -> bool,
	) -> Option<TokenStream> {
		let mut tokens = TokenStream::new();
		while let Some(token) = self.peek() {
			if pred(token) {
				break;
			}
			tokens.append(self.next().unwrap());
		}
		if tokens.is_empty() {
			let span = self.peek().map(|t| t.span());
			self.expected(expected, span);
			return None;
		}
		Some(tokens)
	}
}

/// check if a [`TokenTree`] is a punctuation of symbol `char`
fn is_punct(token: &TokenTree, char: char) -> bool {
	matches!(token, TokenTree::Punct(punct) if punct.as_char() == char)
}

/// main function of `chunk` macro generation
///
/// **syntax:** `
///     ({when(chunked)} -> stream_ident:ident ',')
///     ({when(spanned)} -> span:expr ',')
///     tokens+:tt
/// `
fn chunk_impl(input: TokenStream, chunked: bool, spanned: bool) -> TokenStream {
	let mut cur = Cursor::new(input, Span::call_site());

	cur.add(oquote! { #[allow(unused)] use ::chunked_quote::__private::{
		TokenStream as __TS, TokenStreamExt as _, Delimiter as __Del, ToTokens as __TT,
	}; });

	let stream_ident = if chunked {
		let Some(stream_ident) = cur.eat_ident() else { return cur.res };
		if !cur.eat_punct(',') {
			return cur.res;
		};
		stream_ident
	} else {
		cur.add(oquote! { let mut __stream = __TS::new(); });
		Ident::new("__stream", Span::call_site())
	};

	if spanned {
		let Some(span) = cur.eat_until("an expression", |t| is_punct(t, ',')) else {
			return cur.res;
		};
		if !cur.eat_punct(',') {
			return cur.res;
		}
		cur.add(oquote! { let __span = #span; });
	} else {
		cur.add(oquote! { let __span = ::chunked_quote::__private::Span::call_site(); })
	}

	quote_stream(&mut cur, &stream_ident);
	// `quote` and `quote_spanned` returns the created stream
	if !chunked {
		cur.add(Some(stream_ident));
	}
	let res = cur.res;
	oquote! { { #res } }
}

/// generate [`Punct`] append logic
fn quote_punct(cur: &mut Cursor, punct: &Punct, stream_ident: &Ident) {
	let char = punct.as_char();
	match punct.spacing() {
		Spacing::Joint => {
			cur.add(oquote! { #stream_ident.__push_punct_joint(#char, __span); })
		}
		Spacing::Alone => {
			cur.add(oquote! { #stream_ident.__push_punct_alone(#char, __span); })
		}
	}
}

/// generate [`Ident`] append logic
fn quote_ident(cur: &mut Cursor, ident: &Ident, stream_ident: &Ident) {
	let name = ident.to_string();
	match name.strip_prefix("r#") {
		Some(name) => cur.add(oquote! { #stream_ident.__push_ident_raw(#name, __span); }),
		None => cur.add(oquote! { #stream_ident.__push_ident(#name, __span); }),
	}
}

/// generate [`Literal`] append logic
fn quote_literal(cur: &mut Cursor, lit: &Literal, stream_ident: &Ident) {
	let text = lit.to_string();
	cur.add(oquote! { #stream_ident.__push_literal(#text, __span); });
}

/// generate [`Group`] append logic
fn quote_group(cur: &mut Cursor, group: &Group, stream_ident: &Ident) {
	let delimiter = match group.delimiter() {
		Delimiter::Parenthesis => oquote! { __Del::Parenthesis },
		Delimiter::Brace => oquote! { __Del::Brace },
		Delimiter::Bracket => oquote! { __Del::Bracket },
		Delimiter::None => oquote! { __Del::None },
	};

	let mut inner_cur = Cursor::new(group.stream(), group.span_close());
	quote_stream(&mut inner_cur, stream_ident);
	let inner = inner_cur.res;

	cur.add(oquote! { #stream_ident.__push_group(#delimiter, __span, {
		let mut #stream_ident = __TS::new();
		#inner
		#stream_ident
	}); });
}

/// generate append logic for whole [`TokenStream`]
fn quote_stream(cur: &mut Cursor, stream_ident: &Ident) {
	while let Some(token) = cur.next() {
		match token {
			TokenTree::Ident(ident) => quote_ident(cur, &ident, stream_ident),
			TokenTree::Literal(lit) => quote_literal(cur, &lit, stream_ident),
			TokenTree::Punct(p) if p.as_char() == '#' => {
				handle_directive(cur, stream_ident);
			}
			TokenTree::Punct(punct) => quote_punct(cur, &punct, stream_ident),
			TokenTree::Group(group) => quote_group(cur, &group, stream_ident),
		}
	}
}

/// transform `kw:ident expr:expr '#' body:curly_group` -> `#kw #expr #{ quote_stream(body) }`
fn kw_expr_body(cur: &mut Cursor, ident: Ident, stream_ident: &Ident) -> Option<()> {
	cur.add(Some(ident));

	// we are a utility macro not rustc itself, and this just work.
	let expr = cur.eat_until("an expression", |t| is_punct(t, '#'))?;
	cur.add(expr);

	body(cur, true, stream_ident)
}

/// transform `'#' body:curly_group` -> `quote_stream(body)`
fn body(cur: &mut Cursor, eat_hash: bool, stream_ident: &Ident) -> Option<()> {
	if eat_hash {
		cur.eat_punct('#').then_some(())?;
	}
	let body = cur.eat_brace()?;

	let mut inner_cur = Cursor::new(body.stream(), body.span_close());
	quote_stream(&mut inner_cur, stream_ident);

	cur.add(Some(Group::new(Delimiter::Brace, inner_cur.res)));
	Some(())
}

/// check if a [`TokenTree`] is a brace group
fn is_brace(token: &TokenTree) -> bool {
	matches!(token, TokenTree::Group(group) if group.delimiter() == Delimiter::Brace)
}

/// generate append logic for `# 'match' expr '{' pat => '#' body:curly_group '}'`
fn handle_match(cur: &mut Cursor, ident: Ident, stream_ident: &Ident) -> Option<()> {
	cur.add(Some(ident));
	let expr = cur.eat_until("an expression", is_brace)?;
	cur.add(expr);

	let arms = cur.eat_brace()?;
	let mut arms_cur = Cursor::new(arms.stream(), arms.span_close());

	while let Some(token) = arms_cur.next() {
		if is_punct(&token, '#') {
			body(&mut arms_cur, false, stream_ident)?;
		} else {
			arms_cur.add(Some(token));
		}
	}

	cur.add(Some(Group::new(Delimiter::Brace, arms_cur.res)));
	Some(())
}

/// router of the interpolation logic
fn handle_directive(cur: &mut Cursor, stream_ident: &Ident) -> Option<()> {
	match cur.next() {
		// `'#' '#'` -> `#`
		Some(TokenTree::Punct(p)) if p.as_char() == '#' => {
			quote_punct(cur, &p, stream_ident)
		}
		// `'#' body:curly_group` -> `to_tokens(body)`
		Some(TokenTree::Group(group)) if group.delimiter() == Delimiter::Brace => {
			cur.add(oquote! { __TT::to_tokens(&#group, &mut #stream_ident); })
		}
		Some(TokenTree::Ident(ident)) => match &*ident.to_string() {
			// `'#' op:ident expr:expr '#' body:curly_group`
			"if" | "for" | "while" => kw_expr_body(cur, ident, stream_ident)?,
			"else" => {
				if cur.peek_kw("if") {
					kw_expr_body(cur, ident, stream_ident)?
				} else {
					cur.add(Some(ident));
					body(cur, true, stream_ident)?
				}
			}
			"match" => handle_match(cur, ident, stream_ident)?,
			// `'#' 'do' block:curly_group` -> apend `block` directly
			"do" => {
				let block = cur.eat_brace()?;
				cur.add(oquote! {{#[allow(unused_braces)] #block }})
			}
			// `'#' ident` -> `to_tokens(ident)`
			_ => cur.add(oquote! { __TT::to_tokens(&#ident, &mut #stream_ident); }),
		},
		t => cur.expected("an identifier, `#` or `{`", t.map(|t| t.span())),
	};
	Some(())
}