From 71337f6682c27233d9c7bf9c5e3383da451ba604 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 8 Dec 2023 11:26:22 +0100
Subject: [PATCH] fix: Fix `concat_bytes!` expansion

---
 .../macro_expansion_tests/builtin_fn_macro.rs |  4 +--
 .../macro_expansion_tests/mbe/regression.rs   | 26 +++++++++++++++++++
 crates/hir-expand/src/builtin_fn_macro.rs     | 22 ++++++++++++++--
 crates/mbe/src/token_map.rs                   | 16 +++++++++---
 4 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/crates/hir-def/src/macro_expansion_tests/builtin_fn_macro.rs b/crates/hir-def/src/macro_expansion_tests/builtin_fn_macro.rs
index 106ead83fa..514219ee71 100644
--- a/crates/hir-def/src/macro_expansion_tests/builtin_fn_macro.rs
+++ b/crates/hir-def/src/macro_expansion_tests/builtin_fn_macro.rs
@@ -468,12 +468,12 @@ macro_rules! concat_bytes {}
 
 fn main() { concat_bytes!(b'A', b"BC", [68, b'E', 70]); }
 "##,
-        expect![[r##"
+        expect![[r#"
 #[rustc_builtin_macro]
 macro_rules! concat_bytes {}
 
 fn main() { [b'A', 66, 67, 68, b'E', 70]; }
-"##]],
+"#]],
     );
 }
 
diff --git a/crates/hir-def/src/macro_expansion_tests/mbe/regression.rs b/crates/hir-def/src/macro_expansion_tests/mbe/regression.rs
index 2886b2a366..9010050ee6 100644
--- a/crates/hir-def/src/macro_expansion_tests/mbe/regression.rs
+++ b/crates/hir-def/src/macro_expansion_tests/mbe/regression.rs
@@ -1004,3 +1004,29 @@ fn main() {
 "##]],
     );
 }
+
+#[test]
+fn eager_concat_bytes_panic() {
+    check(
+        r#"
+#[rustc_builtin_macro]
+#[macro_export]
+macro_rules! concat_bytes {}
+
+fn main() {
+    let x = concat_bytes!(2);
+}
+
+"#,
+        expect![[r#"
+#[rustc_builtin_macro]
+#[macro_export]
+macro_rules! concat_bytes {}
+
+fn main() {
+    let x = /* error: unexpected token in input */[];
+}
+
+"#]],
+    );
+}
diff --git a/crates/hir-expand/src/builtin_fn_macro.rs b/crates/hir-expand/src/builtin_fn_macro.rs
index 903c21c84e..c8f04bfee5 100644
--- a/crates/hir-expand/src/builtin_fn_macro.rs
+++ b/crates/hir-expand/src/builtin_fn_macro.rs
@@ -6,6 +6,7 @@ use base_db::{
 };
 use cfg::CfgExpr;
 use either::Either;
+use itertools::Itertools;
 use mbe::{parse_exprs_with_sep, parse_to_token_tree};
 use syntax::{
     ast::{self, AstToken},
@@ -491,8 +492,25 @@ fn concat_bytes_expand(
             }
         }
     }
-    let ident = tt::Ident { text: bytes.join(", ").into(), span };
-    ExpandResult { value: quote!(span =>[#ident]), err }
+    let value = tt::Subtree {
+        delimiter: tt::Delimiter { open: span, close: span, kind: tt::DelimiterKind::Bracket },
+        token_trees: {
+            Itertools::intersperse_with(
+                bytes.into_iter().map(|it| {
+                    tt::TokenTree::Leaf(tt::Leaf::Literal(tt::Literal { text: it.into(), span }))
+                }),
+                || {
+                    tt::TokenTree::Leaf(tt::Leaf::Punct(tt::Punct {
+                        char: ',',
+                        spacing: tt::Spacing::Alone,
+                        span,
+                    }))
+                },
+            )
+            .collect()
+        },
+    };
+    ExpandResult { value, err }
 }
 
 fn concat_bytes_expand_subtree(
diff --git a/crates/mbe/src/token_map.rs b/crates/mbe/src/token_map.rs
index 28b39b4f1e..7d15812f8c 100644
--- a/crates/mbe/src/token_map.rs
+++ b/crates/mbe/src/token_map.rs
@@ -2,7 +2,7 @@
 
 use std::hash::Hash;
 
-use stdx::itertools::Itertools;
+use stdx::{always, itertools::Itertools};
 use syntax::{TextRange, TextSize};
 use tt::Span;
 
@@ -21,13 +21,23 @@ impl<S: Span> SpanMap<S> {
     /// Finalizes the [`SpanMap`], shrinking its backing storage and validating that the offsets are
     /// in order.
     pub fn finish(&mut self) {
-        assert!(self.spans.iter().tuple_windows().all(|(a, b)| a.0 < b.0));
+        always!(
+            self.spans.iter().tuple_windows().all(|(a, b)| a.0 < b.0),
+            "spans are not in order"
+        );
         self.spans.shrink_to_fit();
     }
 
     /// Pushes a new span onto the [`SpanMap`].
     pub fn push(&mut self, offset: TextSize, span: S) {
-        debug_assert!(self.spans.last().map_or(true, |&(last_offset, _)| last_offset < offset));
+        if cfg!(debug_assertions) {
+            if let Some(&(last_offset, _)) = self.spans.last() {
+                assert!(
+                    last_offset < offset,
+                    "last_offset({last_offset:?}) must be smaller than offset({offset:?})"
+                );
+            }
+        }
         self.spans.push((offset, span));
     }