mbe: Refactor the logic for ${concat}

* Now accept numbers, chars
* Suffixes are stripped (needs more testing)
* Report specific locations of errors
* TODO: handle idents the same for expanded tokens

Author: tgross35

compiler/rustc_builtin_macros/src/concat.rs

@@ -24,6 +24,8 @@ pub(crate) fn expand_concat(
     let mut guar = None;
     for e in es {
         match e.kind {
+            // For consistent user experience, please keep this in sync with the handling of
+            // literals in `rustc_expand::mbe::metavar_expr` `${concat()}`!
             ExprKind::Lit(token_lit) => match LitKind::from_token_lit(token_lit) {
                 Ok(LitKind::Str(s, _) | LitKind::Float(s, _)) => {
                     accumulator.push_str(s.as_str());

compiler/rustc_expand/messages.ftl

@@ -133,6 +133,21 @@ expand_module_multiple_candidates =
 expand_must_repeat_once =
     this must repeat at least once
 
+expand_mve_concat_invalid =
+    invalid item within a `{"${concat(...)}"}` expression
+    .expr_ident = expanding this `concat(...)` expression
+    .invalid_ident = this literal produced an invalid identifier
+    .float_lit = float literals cannot be concatenated
+    .c_str_lit = C string literals cannot be concatenated
+    .b_str_lit = byte literals cannot be concatenated
+    .raw_ident = raw identifiers cannot be concatenated
+    .unsupported = unsupported input for `concat(...)`
+    .valid_types = `concat` can join {$valid}
+    .expected_metavar = expected an identifier; got `{$found}`
+    .expected_metavar_dollar = todo
+    .unexpected_group = todo
+    .hi_label = todo
+
 expand_mve_expected_ident =
     expected an identifier
     .not_ident = not a valid identifier

compiler/rustc_expand/src/errors.rs

@@ -540,6 +540,49 @@ mod metavar_exprs {
         pub valid_expr_list: &'static str,
     }
 
+    #[derive(Diagnostic)]
+    #[diag(expand_mve_concat_invalid)]
+    pub(crate) struct MveConcatInvalid {
+        #[primary_span]
+        pub span: Span,
+        #[subdiagnostic]
+        pub reason: MveConcatInvalidReason,
+        #[help(expand_expr_ident)]
+        pub ident_span: Span,
+        pub valid: &'static str,
+    }
+
+    // TODO: can these be labels rather than notes?
+    #[derive(Subdiagnostic)]
+    pub(crate) enum MveConcatInvalidReason {
+        #[note(expand_invalid_ident)]
+        InvalidIdent,
+        #[note(expand_float_lit)]
+        #[help(expand_valid_types)]
+        FloatLit,
+        #[note(expand_c_str_lit)]
+        #[help(expand_valid_types)]
+        CStrLit,
+        #[note(expand_b_str_lit)]
+        #[help(expand_valid_types)]
+        ByteStrLit,
+        #[note(expand_expected_metavar)]
+        #[label(expand_expected_metavar_dollar)]
+        ExpectedMetavarIdent {
+            found: String,
+            #[primary_span]
+            dollar: Span,
+        },
+        #[note(expand_raw_ident)]
+        RawIdentifier,
+        #[note(expand_unsupported)]
+        #[help(expand_valid_types)]
+        UnsupportedInput,
+        #[note(expand_unexpected_group)]
+        UnexpectedGroup,
+        InvalidLiteral,
+    }
+
     #[derive(Diagnostic)]
     #[diag(expand_mve_expected_ident)]
     pub(crate) struct MveExpectedIdent {

compiler/rustc_expand/src/mbe/metavar_expr.rs

@@ -1,16 +1,19 @@
-use rustc_ast::token::{self, Delimiter, IdentIsRaw, Lit, Token, TokenKind};
+use rustc_ast::token::{self, Delimiter, IdentIsRaw, Token, TokenKind};
 use rustc_ast::tokenstream::{TokenStream, TokenStreamIter, TokenTree};
-use rustc_ast::{LitIntType, LitKind};
+use rustc_ast::{self as ast, LitIntType, LitKind};
 use rustc_ast_pretty::pprust;
-use rustc_errors::{Applicability, PResult};
+use rustc_errors::PResult;
+use rustc_lexer::is_id_continue;
 use rustc_macros::{Decodable, Encodable};
+use rustc_session::errors::create_lit_error;
 use rustc_session::parse::ParseSess;
 use rustc_span::{Ident, Span, Symbol};
 
-use crate::errors::{self, MveExpectedIdentContext};
+use crate::errors::{self, MveConcatInvalidReason, MveExpectedIdentContext};
 
 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";
+pub(crate) const VALID_EXPR_CONCAT_TYPES: &str =
+    "metavariables, identifiers, string literals, and integer literals";
 
 /// Argument specification for a metavariable expression
 #[derive(Clone, Copy)]
@@ -190,7 +193,7 @@ fn iter_span(iter: &TokenStreamIter<'_>) -> Option<Span> {
 pub(crate) enum MetaVarExprConcatElem {
     /// Identifier WITHOUT a preceding dollar sign, which means that this identifier should be
     /// interpreted as a literal.
-    Ident(Ident),
+    Ident(String),
     /// For example, a number or a string.
     Literal(Symbol),
     /// Identifier WITH a preceding dollar sign, which means that this identifier should be
@@ -206,30 +209,92 @@ fn parse_concat<'psess>(
     expr_ident_span: Span,
 ) -> PResult<'psess, MetaVarExpr> {
     let mut result = Vec::new();
+    let dcx = psess.dcx();
     loop {
-        let is_var = try_eat_dollar(iter);
-        let token = parse_token(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));
+        let dollar = try_eat_dollar(iter);
+        let Some(tt) = iter.next() else {
+            // May be hit only with the first iteration (peek is otherwise checked at the end).
+            break;
+        };
+
+        let make_err = |reason| {
+            let err = errors::MveConcatInvalid {
+                span: tt.span(),
+                ident_span: expr_ident_span,
+                reason,
+                valid: VALID_EXPR_CONCAT_TYPES,
+            };
+            Err(dcx.create_err(err))
+        };
+
+        let token = match tt {
+            TokenTree::Token(token, _) => token,
+            TokenTree::Delimited(..) => {
+                return make_err(MveConcatInvalidReason::UnexpectedGroup);
+            }
+        };
+
+        let element = if let Some(dollar) = dollar {
+            // Expecting a metavar
+            let Some((ident, _)) = token.ident() else {
+                return make_err(MveConcatInvalidReason::ExpectedMetavarIdent {
+                    found: pprust::token_to_string(token).into_owned(),
+                    dollar,
+                });
+            };
+
+            // Variables get passed untouched
+            MetaVarExprConcatElem::Var(ident)
+        } else if let TokenKind::Literal(lit) = token.kind {
+            // Preprocess with `from_token_lit` to handle unescaping, float / int literal suffix
+            // stripping.
+            //
+            // For consistent user experience, please keep this in sync with the handling of
+            // literals in `rustc_builtin_macros::concat`!
+            let s = match ast::LitKind::from_token_lit(lit.clone()) {
+                Ok(ast::LitKind::Str(s, _)) => s.to_string(),
+                Ok(ast::LitKind::Float(..)) => {
+                    return make_err(MveConcatInvalidReason::FloatLit);
+                }
+                Ok(ast::LitKind::Char(c)) => c.to_string(),
+                Ok(ast::LitKind::Int(i, _)) => i.to_string(),
+                Ok(ast::LitKind::Bool(b)) => b.to_string(),
+                Ok(ast::LitKind::CStr(..)) => return make_err(MveConcatInvalidReason::CStrLit),
+                Ok(ast::LitKind::Byte(..) | ast::LitKind::ByteStr(..)) => {
+                    return make_err(MveConcatInvalidReason::ByteStrLit);
                 }
-                Ok(elem) => MetaVarExprConcatElem::Ident(elem),
+                Ok(ast::LitKind::Err(_guarantee)) => {
+                    // REVIEW: a diagnostic was already emitted, should we just break?
+                    return make_err(MveConcatInvalidReason::InvalidLiteral);
+                }
+                Err(err) => return Err(create_lit_error(psess, err, lit, token.span)),
+            };
+
+            if !s.chars().all(|ch| is_id_continue(ch)) {
+                // Check that all characters are valid in the middle of an identifier. This doesn't
+                // guarantee that the final identifier is valid (we still need to check it later),
+                // but it allows us to catch errors with specific arguments before expansion time;
+                // for example, string literal "foo.bar" gets flagged before the macro is invoked.
+                return make_err(MveConcatInvalidReason::InvalidIdent);
+            }
+
+            MetaVarExprConcatElem::Ident(s)
+        } else if let Some((elem, is_raw)) = token.ident() {
+            if is_raw == IdentIsRaw::Yes {
+                return make_err(MveConcatInvalidReason::RawIdentifier);
             }
+            MetaVarExprConcatElem::Ident(elem.as_str().to_string())
+        } else {
+            return make_err(MveConcatInvalidReason::UnsupportedInput);
         };
+
         result.push(element);
+
         if iter.peek().is_none() {
+            // break before trying to eat the comma
             break;
         }
+
         if !try_eat_comma(iter) {
             return Err(psess.dcx().struct_span_err(outer_span, "expected comma"));
         }
@@ -315,43 +380,6 @@ fn parse_ident<'psess>(
     Ok(elem)
 }
 
-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 TokenStreamIter<'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)
-}
-
 /// Tries to move the iterator forward returning `true` if there is a comma. If not, then the
 /// iterator is not modified and the result is `false`.
 fn try_eat_comma(iter: &mut TokenStreamIter<'_>) -> bool {
@@ -362,14 +390,14 @@ fn try_eat_comma(iter: &mut TokenStreamIter<'_>) -> bool {
     false
 }
 
-/// Tries to move the iterator forward returning `true` if there is a dollar sign. If not, then the
-/// iterator is not modified and the result is `false`.
-fn try_eat_dollar(iter: &mut TokenStreamIter<'_>) -> bool {
-    if let Some(TokenTree::Token(Token { kind: token::Dollar, .. }, _)) = iter.peek() {
+/// Tries to move the iterator forward returning `Some(dollar_span)` if there is a dollar sign. If
+/// not, then the iterator is not modified and the result is `None`.
+fn try_eat_dollar(iter: &mut TokenStreamIter<'_>) -> Option<Span> {
+    if let Some(TokenTree::Token(Token { kind: token::Dollar, span }, _)) = iter.peek() {
         let _ = iter.next();
-        return true;
+        return Some(*span);
     }
-    false
+    None
 }
 
 /// Expects that the next item is a dollar sign.
@@ -378,7 +406,7 @@ fn eat_dollar<'psess>(
     psess: &'psess ParseSess,
     span: Span,
 ) -> PResult<'psess, ()> {
-    if try_eat_dollar(iter) {
+    if try_eat_dollar(iter).is_some() {
         return Ok(());
     }
     Err(psess.dcx().struct_span_err(

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -545,9 +545,10 @@ fn metavar_expr_concat<'tx>(
     let dcx = tscx.psess.dcx();
     let mut concatenated = String::new();
     for element in elements.into_iter() {
-        let symbol = match element {
-            MetaVarExprConcatElem::Ident(elem) => elem.name,
-            MetaVarExprConcatElem::Literal(elem) => *elem,
+        let tmp_sym;
+        let sym_str = match element {
+            MetaVarExprConcatElem::Ident(elem) => elem.as_str(),
+            MetaVarExprConcatElem::Literal(elem) => elem.as_str(),
             MetaVarExprConcatElem::Var(ident) => {
                 match matched_from_ident(dcx, *ident, tscx.interp)? {
                     NamedMatch::MatchedSeq(named_matches) => {
@@ -557,16 +558,20 @@ fn metavar_expr_concat<'tx>(
                         match &named_matches[*curr_idx] {
                             // FIXME(c410-f3r) Nested repetitions are unimplemented
                             MatchedSeq(_) => unimplemented!(),
-                            MatchedSingle(pnr) => extract_symbol_from_pnr(dcx, pnr, ident.span)?,
+                            MatchedSingle(pnr) => {
+                                tmp_sym = extract_symbol_from_pnr(tscx, pnr, ident.span)?;
+                                tmp_sym.as_str()
+                            }
                         }
                     }
                     NamedMatch::MatchedSingle(pnr) => {
-                        extract_symbol_from_pnr(dcx, pnr, ident.span)?
+                        tmp_sym = extract_symbol_from_pnr(tscx, pnr, ident.span)?;
+                        tmp_sym.as_str()
                     }
                 }
             }
         };
-        concatenated.push_str(symbol.as_str());
+        concatenated.push_str(sym_str);
     }
     let symbol = nfc_normalize(&concatenated);
     let concatenated_span = tscx.visited_dspan(dspan);
@@ -900,11 +905,13 @@ fn out_of_bounds_err<'a>(dcx: DiagCtxtHandle<'a>, max: usize, span: Span, ty: &s
 }
 
 /// Extracts an metavariable symbol that can be an identifier, a token tree or a literal.
-fn extract_symbol_from_pnr<'a>(
-    dcx: DiagCtxtHandle<'a>,
+// TODO: use the same logic as for metavar_expr
+fn extract_symbol_from_pnr<'tx>(
+    tscx: &mut TranscrCtx<'tx, '_>,
     pnr: &ParseNtResult,
     span_err: Span,
-) -> PResult<'a, Symbol> {
+) -> PResult<'tx, Symbol> {
+    let dcx = tscx.psess.dcx();
     match pnr {
         ParseNtResult::Ident(nt_ident, is_raw) => {
             if let IdentIsRaw::Yes = is_raw {