Auto merge of rust-lang#92066 - Smittyvb:concat_bytes-repeat, r=nagisa

Support [x; n] expressions in concat_bytes!

Currently trying to use `concat_bytes!` with a repeating array value like `[42; 5]` results in an error:
```
error: expected a byte literal
 --> src/main.rs:3:27
  |
3 |     let x = concat_bytes!([3; 4]);
  |                           ^^^^^^
  |
  = note: only byte literals (like `b"foo"`, `b's'`, and `[3, 4, 5]`) can be passed to `concat_bytes!()`
```

This makes it so repeating array syntax can be used the same way normal arrays can be. The RFC doesn't explicitly mention repeat expressions, but it seems reasonable to allow them as well, since normal arrays are allowed.

It is possible to make the compiler get stuck compiling forever with `concat_bytes!([3; 999999999])`, but I don't think that's much of an issue since you can do that already with `const X: [u8; 999999999] = [3; 999999999];`.

Contributes to rust-lang#87555.

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 positive number");
                 }
             }
             ast::ExprKind::Lit(ref lit) => match lit.kind {

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 positive number
+    concat_bytes!([3; -2]); //~ ERROR repeat count is not a positive number
+    concat_bytes!([pie; -2]); //~ ERROR repeat count is not a positive 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 positive number
+    concat_bytes!([[1, 2, 3]; 3]); //~ ERROR cannot concatenate doubly nested array
+    concat_bytes!([[42; 2]; 3]); //~ ERROR cannot concatenate doubly nested array
 }

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 positive number
+  --> $DIR/concat-bytes-error.rs:42:23
+   |
+LL |     concat_bytes!([3; ()]);
+   |                       ^^
+
+error: repeat count is not a positive number
+  --> $DIR/concat-bytes-error.rs:43:23
+   |
+LL |     concat_bytes!([3; -2]);
+   |                       ^^
+
+error: repeat count is not a positive 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 positive 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
 

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",
+    );
 }