Allow concat in repetitions

Author: c410-f3r

compiler/rustc_expand/src/mbe/macro_check.rs

@@ -352,10 +352,10 @@ fn check_occurrences(
             check_ops_is_prefix(psess, node_id, macros, binders, ops, span, name);
         }
         TokenTree::MetaVarExpr(dl, ref mve) => {
-            let Some(name) = mve.ident().map(MacroRulesNormalizedIdent::new) else {
-                return;
-            };
-            check_ops_is_prefix(psess, node_id, macros, binders, ops, dl.entire(), name);
+            mve.metavars((), |_, ident| {
+                let name = MacroRulesNormalizedIdent::new(*ident);
+                check_ops_is_prefix(psess, node_id, macros, binders, ops, dl.entire(), name);
+            });
         }
         TokenTree::Delimited(.., ref del) => {
             check_nested_occurrences(psess, node_id, &del.tts, macros, binders, ops, guar);

compiler/rustc_expand/src/mbe/metavar_expr.rs

@@ -111,10 +111,18 @@ impl MetaVarExpr {
         Ok(rslt)
     }
 
-    pub(crate) fn ident(&self) -> Option<Ident> {
-        match *self {
-            MetaVarExpr::Count(ident, _) | MetaVarExpr::Ignore(ident) => Some(ident),
-            MetaVarExpr::Concat { .. } | MetaVarExpr::Index(..) | MetaVarExpr::Len(..) => None,
+    pub(crate) fn metavars<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,
         }
     }
 }

compiler/rustc_expand/src/mbe/transcribe.rs

@@ -557,17 +557,13 @@ fn lockstep_iter_size(
             }
         }
         TokenTree::MetaVarExpr(_, expr) => {
-            let default_rslt = LockstepIterSize::Unconstrained;
-            let Some(ident) = expr.ident() else {
-                return default_rslt;
-            };
-            let name = MacroRulesNormalizedIdent::new(ident);
-            match lookup_cur_matched(name, interpolations, repeats) {
-                Some(MatchedSeq(ads)) => {
-                    default_rslt.with(LockstepIterSize::Constraint(ads.len(), name))
-                }
-                _ => default_rslt,
-            }
+            expr.metavars(LockstepIterSize::Unconstrained, |lis, ident| {
+                lis.with(lockstep_iter_size(
+                    &TokenTree::MetaVar(ident.span, *ident),
+                    interpolations,
+                    repeats,
+                ))
+            })
         }
         TokenTree::Token(..) => LockstepIterSize::Unconstrained,
     }
@@ -695,7 +691,22 @@ fn transcribe_metavar_expr<'a>(
                 let symbol = match element {
                     MetaVarExprConcatElem::Ident(elem) => elem.name,
                     MetaVarExprConcatElem::Literal(elem) => *elem,
-                    MetaVarExprConcatElem::Var(elem) => extract_var_symbol(dcx, *elem, interp)?,
+                    MetaVarExprConcatElem::Var(ident) => {
+                        match matched_from_ident(dcx, *ident, interp)? {
+                            NamedMatch::MatchedSeq(named_matches) => {
+                                let curr_idx = repeats.last().unwrap().0;
+                                match &named_matches[curr_idx] {
+                                    MatchedSeq(_) => unimplemented!(),
+                                    MatchedSingle(pnr) => {
+                                        extract_symbol_from_pnr(dcx, *ident, pnr)?
+                                    }
+                                }
+                            }
+                            NamedMatch::MatchedSingle(pnr) => {
+                                extract_symbol_from_pnr(dcx, *ident, pnr)?
+                            }
+                        }
+                    }
                 };
                 concatenated.push_str(symbol.as_str());
             }
@@ -752,39 +763,37 @@ fn transcribe_metavar_expr<'a>(
 }
 
 /// Extracts an metavariable symbol that can be an identifier, a token tree or a literal.
-fn extract_var_symbol<'a>(
+fn extract_symbol_from_pnr<'a>(
     dcx: DiagCtxtHandle<'a>,
     ident: Ident,
-    interp: &FxHashMap<MacroRulesNormalizedIdent, NamedMatch>,
+    pnr: &ParseNtResult,
 ) -> PResult<'a, Symbol> {
-    if let NamedMatch::MatchedSingle(pnr) = matched_from_ident(dcx, ident, interp)? {
-        if let ParseNtResult::Ident(nt_ident, is_raw) = pnr {
-            if let IdentIsRaw::Yes = is_raw {
-                return Err(dcx.struct_span_err(ident.span, RAW_IDENT_ERR));
-            }
-            return Ok(nt_ident.name);
+    if let ParseNtResult::Ident(nt_ident, is_raw) = pnr {
+        if let IdentIsRaw::Yes = is_raw {
+            return Err(dcx.struct_span_err(ident.span, RAW_IDENT_ERR));
         }
+        return Ok(nt_ident.name);
+    }
 
-        if let ParseNtResult::Tt(TokenTree::Token(Token { kind, .. }, _)) = pnr {
-            if let TokenKind::Ident(symbol, is_raw) = kind {
-                if let IdentIsRaw::Yes = is_raw {
-                    return Err(dcx.struct_span_err(ident.span, RAW_IDENT_ERR));
-                }
-                return Ok(*symbol);
-            }
-
-            if let TokenKind::Literal(Lit { kind: LitKind::Str, symbol, suffix: None }) = kind {
-                return Ok(*symbol);
+    if let ParseNtResult::Tt(TokenTree::Token(Token { kind, .. }, _)) = pnr {
+        if let TokenKind::Ident(symbol, is_raw) = kind {
+            if let IdentIsRaw::Yes = is_raw {
+                return Err(dcx.struct_span_err(ident.span, RAW_IDENT_ERR));
             }
+            return Ok(*symbol);
         }
 
-        if let ParseNtResult::Nt(nt) = pnr
-            && let Nonterminal::NtLiteral(expr) = &**nt
-            && let ExprKind::Lit(Lit { kind: LitKind::Str, symbol, suffix: None }) = &expr.kind
-        {
+        if let TokenKind::Literal(Lit { kind: LitKind::Str, symbol, suffix: None }) = kind {
             return Ok(*symbol);
         }
     }
+
+    if let ParseNtResult::Nt(nt) = pnr
+        && let Nonterminal::NtLiteral(expr) = &**nt
+        && let ExprKind::Lit(Lit { kind: LitKind::Str, symbol, suffix: None }) = &expr.kind
+    {
+        return Ok(*symbol);
+    }
     Err(dcx
         .struct_err("metavariables of `${concat(..)}` must be of type `ident`, `literal` or `tt`")
         .with_note("currently only string literals are supported")

tests/ui/macros/macro-metavar-expr-concat/repetitions.rs

@@ -0,0 +1,18 @@
+//@ run-pass
+
+#![feature(macro_metavar_expr_concat)]
+
+macro_rules! one_rep {
+    ( $($a:ident)* ) => {
+        $(
+            const ${concat($a, Z)}: i32 = 3;
+        )*
+    };
+}
+
+fn main() {
+    one_rep!(A B C);
+    assert_eq!(AZ, 3);
+    assert_eq!(BZ, 3);
+    assert_eq!(CZ, 3);
+}