Support [x; n] expressions in concat_bytes!

Contributes to #87555.

4 files changed, 135 insertions, 42 deletions

diff --git a/compiler/rustc_builtin_macros/src/concat_bytes.rs b/compiler/rustc_builtin_macros/src/concat_bytes.rs
index a107f5993b5..87920e51752 100644
--- a/compiler/rustc_builtin_macros/src/concat_bytes.rs
+++ b/compiler/rustc_builtin_macros/src/concat_bytes.rs
@@ -72,6 +72,52 @@ fn invalid_type_err(cx: &mut base::ExtCtxt<'_>, expr: &P<rustc_ast::Expr>, is_ne
     }
 }
 
+fn handle_array_element(
+    cx: &mut base::ExtCtxt<'_>,
+    has_errors: &mut bool,
+    missing_literals: &mut Vec<rustc_span::Span>,
+    expr: &P<rustc_ast::Expr>,
+) -> Option<u8> {
+    match expr.kind {
+        ast::ExprKind::Array(_) | ast::ExprKind::Repeat(_, _) => {
+            if !*has_errors {
+                cx.span_err(expr.span, "cannot concatenate doubly nested array");
+            }
+            *has_errors = true;
+            None
+        }
+        ast::ExprKind::Lit(ref lit) => match lit.kind {
+            ast::LitKind::Int(
+                val,
+                ast::LitIntType::Unsuffixed | ast::LitIntType::Unsigned(ast::UintTy::U8),
+            ) if val <= u8::MAX.into() => Some(val as u8),
+
+            ast::LitKind::Byte(val) => Some(val),
+            ast::LitKind::ByteStr(_) => {
+                if !*has_errors {
+                    cx.struct_span_err(expr.span, "cannot concatenate doubly nested array")
+                        .note("byte strings are treated as arrays of bytes")
+                        .help("try flattening the array")
+                        .emit();
+                }
+                *has_errors = true;
+                None
+            }
+            _ => {
+                if !*has_errors {
+                    invalid_type_err(cx, expr, true);
+                }
+                *has_errors = true;
+                None
+            }
+        },
+        _ => {
+            missing_literals.push(expr.span);
+            None
+        }
+    }
+}
+
 pub fn expand_concat_bytes(
     cx: &mut base::ExtCtxt<'_>,
     sp: rustc_span::Span,
@@ -88,48 +134,27 @@ pub fn expand_concat_bytes(
         match e.kind {
             ast::ExprKind::Array(ref exprs) => {
                 for expr in exprs {
-                    match expr.kind {
-                        ast::ExprKind::Array(_) => {
-                            if !has_errors {
-                                cx.span_err(expr.span, "cannot concatenate doubly nested array");
-                            }
-                            has_errors = true;
-                        }
-                        ast::ExprKind::Lit(ref lit) => match lit.kind {
-                            ast::LitKind::Int(
-                                val,
-                                ast::LitIntType::Unsuffixed
-                                | ast::LitIntType::Unsigned(ast::UintTy::U8),
-                            ) if val <= u8::MAX.into() => {
-                                accumulator.push(val as u8);
-                            }
-
-                            ast::LitKind::Byte(val) => {
-                                accumulator.push(val);
-                            }
-                            ast::LitKind::ByteStr(_) => {
-                                if !has_errors {
-                                    cx.struct_span_err(
-                                        expr.span,
-                                        "cannot concatenate doubly nested array",
-                                    )
-                                    .note("byte strings are treated as arrays of bytes")
-                                    .help("try flattening the array")
-                                    .emit();
-                                }
-                                has_errors = true;
-                            }
-                            _ => {
-                                if !has_errors {
-                                    invalid_type_err(cx, expr, true);
-                                }
-                                has_errors = true;
-                            }
-                        },
-                        _ => {
-                            missing_literals.push(expr.span);
+                    if let Some(elem) =
+                        handle_array_element(cx, &mut has_errors, &mut missing_literals, expr)
+                    {
+                        accumulator.push(elem);
+                    }
+                }
+            }
+            ast::ExprKind::Repeat(ref expr, ref count) => {
+                if let ast::ExprKind::Lit(ast::Lit {
+                    kind: ast::LitKind::Int(count_val, _), ..
+                }) = count.value.kind
+                {
+                    if let Some(elem) =
+                        handle_array_element(cx, &mut has_errors, &mut missing_literals, expr)
+                    {
+                        for _ in 0..count_val {
+                            accumulator.push(elem);
                         }
                     }
+                } else {
+                    cx.span_err(count.value.span, "repeat count is not a number");
                 }
             }
             ast::ExprKind::Lit(ref lit) => match lit.kind {
diff --git a/src/test/ui/macros/concat-bytes-error.rs b/src/test/ui/macros/concat-bytes-error.rs
index 9b4a9c2cf81..2c7997d5e59 100644
--- a/src/test/ui/macros/concat-bytes-error.rs
+++ b/src/test/ui/macros/concat-bytes-error.rs
@@ -39,4 +39,12 @@ fn main() {
     ]);
     concat_bytes!(5u16); //~ ERROR cannot concatenate numeric literals
     concat_bytes!([5u16]); //~ ERROR numeric literal is not a `u8`
+    concat_bytes!([3; ()]); //~ ERROR repeat count is not a number
+    concat_bytes!([3; -2]); //~ ERROR repeat count is not a number
+    concat_bytes!([pie; -2]); //~ ERROR repeat count is not a number
+    concat_bytes!([pie; 2]); //~ ERROR expected a byte literal
+    concat_bytes!([2.2; 0]); //~ ERROR cannot concatenate float literals
+    concat_bytes!([5.5; ()]); //~ ERROR repeat count is not a number
+    concat_bytes!([[1, 2, 3]; 3]); //~ ERROR cannot concatenate doubly nested array
+    concat_bytes!([[42; 2]; 3]); //~ ERROR cannot concatenate doubly nested array
 }
diff --git a/src/test/ui/macros/concat-bytes-error.stderr b/src/test/ui/macros/concat-bytes-error.stderr
index 1fc2d5c4843..f0c53839b3a 100644
--- a/src/test/ui/macros/concat-bytes-error.stderr
+++ b/src/test/ui/macros/concat-bytes-error.stderr
@@ -127,5 +127,55 @@ error: numeric literal is not a `u8`
 LL |     concat_bytes!([5u16]);
    |                    ^^^^
 
-error: aborting due to 20 previous errors
+error: repeat count is not a number
+  --> $DIR/concat-bytes-error.rs:42:23
+   |
+LL |     concat_bytes!([3; ()]);
+   |                       ^^
+
+error: repeat count is not a number
+  --> $DIR/concat-bytes-error.rs:43:23
+   |
+LL |     concat_bytes!([3; -2]);
+   |                       ^^
+
+error: repeat count is not a number
+  --> $DIR/concat-bytes-error.rs:44:25
+   |
+LL |     concat_bytes!([pie; -2]);
+   |                         ^^
+
+error: expected a byte literal
+  --> $DIR/concat-bytes-error.rs:45:20
+   |
+LL |     concat_bytes!([pie; 2]);
+   |                    ^^^
+   |
+   = note: only byte literals (like `b"foo"`, `b's'`, and `[3, 4, 5]`) can be passed to `concat_bytes!()`
+
+error: cannot concatenate float literals
+  --> $DIR/concat-bytes-error.rs:46:20
+   |
+LL |     concat_bytes!([2.2; 0]);
+   |                    ^^^
+
+error: repeat count is not a number
+  --> $DIR/concat-bytes-error.rs:47:25
+   |
+LL |     concat_bytes!([5.5; ()]);
+   |                         ^^
+
+error: cannot concatenate doubly nested array
+  --> $DIR/concat-bytes-error.rs:48:20
+   |
+LL |     concat_bytes!([[1, 2, 3]; 3]);
+   |                    ^^^^^^^^^
+
+error: cannot concatenate doubly nested array
+  --> $DIR/concat-bytes-error.rs:49:20
+   |
+LL |     concat_bytes!([[42; 2]; 3]);
+   |                    ^^^^^^^
+
+error: aborting due to 28 previous errors
 
diff --git a/src/test/ui/macros/concat-bytes.rs b/src/test/ui/macros/concat-bytes.rs
index 5415cf3fe22..fd8f99417ec 100644
--- a/src/test/ui/macros/concat-bytes.rs
+++ b/src/test/ui/macros/concat-bytes.rs
@@ -3,5 +3,15 @@
 
 fn main() {
     assert_eq!(concat_bytes!(), &[]);
-    assert_eq!(concat_bytes!(b'A', b"BC", [68, b'E', 70]), b"ABCDEF");
+    assert_eq!(
+        concat_bytes!(b'A', b"BC", [68, b'E', 70], [b'G'; 1], [72; 2], [73u8; 3], [65; 0]),
+        b"ABCDEFGHHIII",
+    );
+    assert_eq!(
+        concat_bytes!(
+            concat_bytes!(b"AB", b"CD"),
+            concat_bytes!(b"EF", b"GH"),
+        ),
+        b"ABCDEFGH",
+    );
 }