use rustc_ast::token::{self, Delimiter, IdentIsRaw, Lit, Token, TokenKind};
use rustc_ast::tokenstream::{RefTokenTreeCursor, TokenStream, TokenTree};
use rustc_ast::{LitIntType, LitKind};
use rustc_ast_pretty::pprust;
use rustc_errors::{Applicability, PResult};
use rustc_macros::{Decodable, Encodable};
use rustc_session::parse::ParseSess;
use rustc_span::symbol::Ident;
use rustc_span::{Span, Symbol};
pub(crate) const RAW_IDENT_ERR: &str = "`${concat(..)}` currently does not support raw identifiers";
pub(crate) const UNSUPPORTED_CONCAT_ELEM_ERR: &str = "expected identifier or string literal";
#[derive(Debug, PartialEq, Encodable, Decodable)]
pub(crate) enum MetaVarExpr {
Concat(Box<[MetaVarExprConcatElem]>),
Count(Ident, usize),
Ignore(Ident),
Index(usize),
Len(usize),
}
impl MetaVarExpr {
pub(crate) fn parse<'psess>(
input: &TokenStream,
outer_span: Span,
psess: &'psess ParseSess,
) -> PResult<'psess, MetaVarExpr> {
let mut tts = input.trees();
let ident = parse_ident(&mut tts, psess, outer_span)?;
let Some(TokenTree::Delimited(.., Delimiter::Parenthesis, args)) = tts.next() else {
let msg = "meta-variable expression parameter must be wrapped in parentheses";
return Err(psess.dcx().struct_span_err(ident.span, msg));
};
check_trailing_token(&mut tts, psess)?;
let mut iter = args.trees();
let rslt = match ident.as_str() {
"concat" => {
let mut result = Vec::new();
loop {
let is_var = try_eat_dollar(&mut iter);
let token = parse_token(&mut iter, psess, outer_span)?;
let element = if is_var {
MetaVarExprConcatElem::Var(parse_ident_from_token(psess, token)?)
} else if let TokenKind::Literal(Lit {
kind: token::LitKind::Str,
symbol,
suffix: None,
}) = token.kind
{
MetaVarExprConcatElem::Literal(symbol)
} else {
match parse_ident_from_token(psess, token) {
Err(err) => {
err.cancel();
return Err(psess
.dcx()
.struct_span_err(token.span, UNSUPPORTED_CONCAT_ELEM_ERR));
}
Ok(elem) => MetaVarExprConcatElem::Ident(elem),
}
};
result.push(element);
if iter.look_ahead(0).is_none() {
break;
}
if !try_eat_comma(&mut iter) {
return Err(psess.dcx().struct_span_err(outer_span, "expected comma"));
}
}
if result.len() < 2 {
return Err(psess
.dcx()
.struct_span_err(ident.span, "`concat` must have at least two elements"));
}
MetaVarExpr::Concat(result.into())
}
"count" => parse_count(&mut iter, psess, ident.span)?,
"ignore" => {
eat_dollar(&mut iter, psess, ident.span)?;
MetaVarExpr::Ignore(parse_ident(&mut iter, psess, ident.span)?)
}
"index" => MetaVarExpr::Index(parse_depth(&mut iter, psess, ident.span)?),
"len" => MetaVarExpr::Len(parse_depth(&mut iter, psess, ident.span)?),
_ => {
let err_msg = "unrecognized meta-variable expression";
let mut err = psess.dcx().struct_span_err(ident.span, err_msg);
err.span_suggestion(
ident.span,
"supported expressions are count, ignore, index and len",
"",
Applicability::MachineApplicable,
);
return Err(err);
}
};
check_trailing_token(&mut iter, psess)?;
Ok(rslt)
}
pub(crate) fn for_each_metavar<A>(&self, mut aux: A, mut cb: impl FnMut(A, &Ident) -> A) -> A {
match self {
MetaVarExpr::Concat(elems) => {
for elem in elems {
if let MetaVarExprConcatElem::Var(ident) = elem {
aux = cb(aux, ident)
}
}
aux
}
MetaVarExpr::Count(ident, _) | MetaVarExpr::Ignore(ident) => cb(aux, ident),
MetaVarExpr::Index(..) | MetaVarExpr::Len(..) => aux,
}
}
}
#[derive(Debug, Decodable, Encodable, PartialEq)]
pub(crate) enum MetaVarExprConcatElem {
Ident(Ident),
Literal(Symbol),
Var(Ident),
}
fn check_trailing_token<'psess>(
iter: &mut RefTokenTreeCursor<'_>,
psess: &'psess ParseSess,
) -> PResult<'psess, ()> {
if let Some(tt) = iter.next() {
let mut diag = psess
.dcx()
.struct_span_err(tt.span(), format!("unexpected token: {}", pprust::tt_to_string(tt)));
diag.span_note(tt.span(), "meta-variable expression must not have trailing tokens");
Err(diag)
} else {
Ok(())
}
}
fn parse_count<'psess>(
iter: &mut RefTokenTreeCursor<'_>,
psess: &'psess ParseSess,
span: Span,
) -> PResult<'psess, MetaVarExpr> {
eat_dollar(iter, psess, span)?;
let ident = parse_ident(iter, psess, span)?;
let depth = if try_eat_comma(iter) {
if iter.look_ahead(0).is_none() {
return Err(psess.dcx().struct_span_err(
span,
"`count` followed by a comma must have an associated index indicating its depth",
));
}
parse_depth(iter, psess, span)?
} else {
0
};
Ok(MetaVarExpr::Count(ident, depth))
}
fn parse_depth<'psess>(
iter: &mut RefTokenTreeCursor<'_>,
psess: &'psess ParseSess,
span: Span,
) -> PResult<'psess, usize> {
let Some(tt) = iter.next() else { return Ok(0) };
let TokenTree::Token(Token { kind: TokenKind::Literal(lit), .. }, _) = tt else {
return Err(psess
.dcx()
.struct_span_err(span, "meta-variable expression depth must be a literal"));
};
if let Ok(lit_kind) = LitKind::from_token_lit(*lit)
&& let LitKind::Int(n_u128, LitIntType::Unsuffixed) = lit_kind
&& let Ok(n_usize) = usize::try_from(n_u128.get())
{
Ok(n_usize)
} else {
let msg = "only unsuffixes integer literals are supported in meta-variable expressions";
Err(psess.dcx().struct_span_err(span, msg))
}
}
fn parse_ident<'psess>(
iter: &mut RefTokenTreeCursor<'_>,
psess: &'psess ParseSess,
fallback_span: Span,
) -> PResult<'psess, Ident> {
let token = parse_token(iter, psess, fallback_span)?;
parse_ident_from_token(psess, token)
}
fn parse_ident_from_token<'psess>(
psess: &'psess ParseSess,
token: &Token,
) -> PResult<'psess, Ident> {
if let Some((elem, is_raw)) = token.ident() {
if let IdentIsRaw::Yes = is_raw {
return Err(psess.dcx().struct_span_err(elem.span, RAW_IDENT_ERR));
}
return Ok(elem);
}
let token_str = pprust::token_to_string(token);
let mut err = psess
.dcx()
.struct_span_err(token.span, format!("expected identifier, found `{token_str}`"));
err.span_suggestion(
token.span,
format!("try removing `{token_str}`"),
"",
Applicability::MaybeIncorrect,
);
Err(err)
}
fn parse_token<'psess, 't>(
iter: &mut RefTokenTreeCursor<'t>,
psess: &'psess ParseSess,
fallback_span: Span,
) -> PResult<'psess, &'t Token> {
let Some(tt) = iter.next() else {
return Err(psess.dcx().struct_span_err(fallback_span, UNSUPPORTED_CONCAT_ELEM_ERR));
};
let TokenTree::Token(token, _) = tt else {
return Err(psess.dcx().struct_span_err(tt.span(), UNSUPPORTED_CONCAT_ELEM_ERR));
};
Ok(token)
}
fn try_eat_comma(iter: &mut RefTokenTreeCursor<'_>) -> bool {
if let Some(TokenTree::Token(Token { kind: token::Comma, .. }, _)) = iter.look_ahead(0) {
let _ = iter.next();
return true;
}
false
}
fn try_eat_dollar(iter: &mut RefTokenTreeCursor<'_>) -> bool {
if let Some(TokenTree::Token(Token { kind: token::Dollar, .. }, _)) = iter.look_ahead(0) {
let _ = iter.next();
return true;
}
false
}
fn eat_dollar<'psess>(
iter: &mut RefTokenTreeCursor<'_>,
psess: &'psess ParseSess,
span: Span,
) -> PResult<'psess, ()> {
if let Some(TokenTree::Token(Token { kind: token::Dollar, .. }, _)) = iter.look_ahead(0) {
let _ = iter.next();
return Ok(());
}
Err(psess.dcx().struct_span_err(
span,
"meta-variables within meta-variable expressions must be referenced using a dollar sign",
))
}