From 494fab202a7ec00f79636dcece8ca5e2742efff0 Mon Sep 17 00:00:00 2001
From: ihciah <ihciah@gmail.com>
Date: Fri, 25 Feb 2022 18:46:11 +0800
Subject: [PATCH] feat: support concat_bytes

---
 .../macro_expansion_tests/builtin_fn_macro.rs | 18 +++++
 crates/hir_expand/src/builtin_fn_macro.rs     | 75 +++++++++++++++++++
 crates/hir_expand/src/name.rs                 |  1 +
 3 files changed, 94 insertions(+)

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 6e91301ecd..919dd6c07f 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
@@ -313,6 +313,24 @@ fn main() { "foor0bar\nfalse"; }
     );
 }
 
+#[test]
+fn test_concat_bytes_expand() {
+    check(
+        r##"
+#[rustc_builtin_macro]
+macro_rules! concat_bytes {}
+
+fn main() { concat_bytes!(b'A', b"BC", [68, b'E', 70]); }
+"##,
+        expect![[r##"
+#[rustc_builtin_macro]
+macro_rules! concat_bytes {}
+
+fn main() { [b'A', 66, 67, 68, b'E', 70]; }
+"##]],
+    );
+}
+
 #[test]
 fn test_concat_with_captured_expr() {
     check(
diff --git a/crates/hir_expand/src/builtin_fn_macro.rs b/crates/hir_expand/src/builtin_fn_macro.rs
index 5876be81b4..bf6cc9989a 100644
--- a/crates/hir_expand/src/builtin_fn_macro.rs
+++ b/crates/hir_expand/src/builtin_fn_macro.rs
@@ -121,6 +121,7 @@ register_builtin! {
     (compile_error, CompileError) => compile_error_expand,
     (concat, Concat) => concat_expand,
     (concat_idents, ConcatIdents) => concat_idents_expand,
+    (concat_bytes, ConcatBytes) => concat_bytes_expand,
     (include, Include) => include_expand,
     (include_bytes, IncludeBytes) => include_bytes_expand,
     (include_str, IncludeStr) => include_str_expand,
@@ -359,6 +360,12 @@ fn unquote_str(lit: &tt::Literal) -> Option<String> {
     token.value().map(|it| it.into_owned())
 }
 
+fn unquote_byte_string(lit: &tt::Literal) -> Option<Vec<u8>> {
+    let lit = ast::make::tokens::literal(&lit.to_string());
+    let token = ast::ByteString::cast(lit)?;
+    token.value().map(|it| it.into_owned())
+}
+
 fn compile_error_expand(
     _db: &dyn AstDatabase,
     _id: MacroCallId,
@@ -422,6 +429,74 @@ fn concat_expand(
     ExpandResult { value: ExpandedEager::new(quote!(#text)), err }
 }
 
+fn concat_bytes_expand(
+    _db: &dyn AstDatabase,
+    _arg_id: MacroCallId,
+    tt: &tt::Subtree,
+) -> ExpandResult<ExpandedEager> {
+    let mut bytes = Vec::new();
+    let mut err = None;
+    for (i, t) in tt.token_trees.iter().enumerate() {
+        match t {
+            tt::TokenTree::Leaf(tt::Leaf::Literal(lit)) => {
+                let token = ast::make::tokens::literal(&lit.to_string());
+                match token.kind() {
+                    syntax::SyntaxKind::BYTE => bytes.push(token.text().to_string()),
+                    syntax::SyntaxKind::BYTE_STRING => {
+                        let components = unquote_byte_string(lit).unwrap_or_else(|| Vec::new());
+                        components.into_iter().for_each(|x| bytes.push(x.to_string()));
+                    }
+                    _ => {
+                        err.get_or_insert(mbe::ExpandError::UnexpectedToken.into());
+                        break;
+                    }
+                }
+            }
+            tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) if i % 2 == 1 && punct.char == ',' => (),
+            tt::TokenTree::Subtree(tree)
+                if tree.delimiter_kind() == Some(tt::DelimiterKind::Bracket) =>
+            {
+                if let Err(e) = concat_bytes_expand_subtree(tree, &mut bytes) {
+                    err.get_or_insert(e);
+                    break;
+                }
+            }
+            _ => {
+                err.get_or_insert(mbe::ExpandError::UnexpectedToken.into());
+                break;
+            }
+        }
+    }
+    let ident = tt::Ident { text: bytes.join(", ").into(), id: tt::TokenId::unspecified() };
+    ExpandResult { value: ExpandedEager::new(quote!([#ident])), err }
+}
+
+fn concat_bytes_expand_subtree(
+    tree: &tt::Subtree,
+    bytes: &mut Vec<String>,
+) -> Result<(), ExpandError> {
+    for (ti, tt) in tree.token_trees.iter().enumerate() {
+        match tt {
+            tt::TokenTree::Leaf(tt::Leaf::Literal(lit)) => {
+                let lit = ast::make::tokens::literal(&lit.to_string());
+                match lit.kind() {
+                    syntax::SyntaxKind::BYTE | syntax::SyntaxKind::INT_NUMBER => {
+                        bytes.push(lit.text().to_string())
+                    }
+                    _ => {
+                        return Err(mbe::ExpandError::UnexpectedToken.into());
+                    }
+                }
+            }
+            tt::TokenTree::Leaf(tt::Leaf::Punct(punct)) if ti % 2 == 1 && punct.char == ',' => (),
+            _ => {
+                return Err(mbe::ExpandError::UnexpectedToken.into());
+            }
+        }
+    }
+    Ok(())
+}
+
 fn concat_idents_expand(
     _db: &dyn AstDatabase,
     _arg_id: MacroCallId,
diff --git a/crates/hir_expand/src/name.rs b/crates/hir_expand/src/name.rs
index c36bd09e2b..7deec253da 100644
--- a/crates/hir_expand/src/name.rs
+++ b/crates/hir_expand/src/name.rs
@@ -233,6 +233,7 @@ pub mod known {
         column,
         compile_error,
         concat_idents,
+        concat_bytes,
         concat,
         const_format_args,
         core_panic,