Teach concat!() to concatenate byte str literals

Author: estebank

src/libsyntax/ast.rs

@@ -1300,6 +1300,11 @@ impl LitKind {
         }
     }
 
+    /// Returns a `LitKind::ByteStr` from `content`.
+    pub fn new_byte_str(content: Vec<u8>) -> LitKind {
+        LitKind::ByteStr(Lrc::new(content))
+    }
+
     /// Returns true if this is a numeric literal.
     pub fn is_numeric(&self) -> bool {
         match *self {

src/libsyntax_ext/concat.rs

@@ -13,42 +13,54 @@ use syntax::ext::base;
 use syntax::ext::build::AstBuilder;
 use syntax::symbol::Symbol;
 use syntax::tokenstream;
-use syntax_pos;
+use syntax_pos::Span;
 
 use std::string::String;
 
 pub fn expand_syntax_ext(
     cx: &mut base::ExtCtxt,
-    sp: syntax_pos::Span,
+    sp: Span,
     tts: &[tokenstream::TokenTree],
 ) -> Box<dyn base::MacResult + 'static> {
     let es = match base::get_exprs_from_tts(cx, sp, tts) {
         Some(e) => e,
         None => return base::DummyResult::expr(sp),
     };
-    let mut accumulator = String::new();
+    let mut string_accumulator = String::new();
+    let mut string_pos = vec![];
+    let mut b_accumulator: Vec<u8> = vec![];
+    let mut b_pos: Vec<Span> = vec![];
     let mut missing_literal = vec![];
     for e in es {
         match e.node {
             ast::ExprKind::Lit(ref lit) => match lit.node {
                 ast::LitKind::Str(ref s, _)
                 | ast::LitKind::Float(ref s, _)
                 | ast::LitKind::FloatUnsuffixed(ref s) => {
-                    accumulator.push_str(&s.as_str());
+                    string_accumulator.push_str(&s.as_str());
+                    string_pos.push(e.span);
                 }
                 ast::LitKind::Char(c) => {
-                    accumulator.push(c);
+                    string_accumulator.push(c);
+                    string_pos.push(e.span);
                 }
                 ast::LitKind::Int(i, ast::LitIntType::Unsigned(_))
                 | ast::LitKind::Int(i, ast::LitIntType::Signed(_))
                 | ast::LitKind::Int(i, ast::LitIntType::Unsuffixed) => {
-                    accumulator.push_str(&i.to_string());
+                    string_accumulator.push_str(&i.to_string());
+                    string_pos.push(e.span);
                 }
                 ast::LitKind::Bool(b) => {
-                    accumulator.push_str(&b.to_string());
+                    string_accumulator.push_str(&b.to_string());
+                    string_pos.push(e.span);
                 }
-                ast::LitKind::Byte(..) | ast::LitKind::ByteStr(..) => {
-                    cx.span_err(e.span, "cannot concatenate a byte string literal");
+                ast::LitKind::Byte(byte) => {
+                    b_accumulator.push(byte);
+                    b_pos.push(e.span);
+                }
+                ast::LitKind::ByteStr(ref b_str) => {
+                    b_accumulator.extend(b_str.iter());
+                    b_pos.push(e.span);
                 }
             },
             _ => {
@@ -61,6 +73,36 @@ pub fn expand_syntax_ext(
         err.note("only literals (like `\"foo\"`, `42` and `3.14`) can be passed to `concat!()`");
         err.emit();
     }
+    // Do not allow mixing "" and b""
+    if string_accumulator.len() > 0 && b_accumulator.len() > 0 {
+        let mut err = cx.struct_span_err(
+            b_pos.clone(),
+            "cannot concatenate a byte string literal with string literals",
+        );
+        for pos in &b_pos {
+            err.span_label(*pos, "byte string literal");
+        }
+        for pos in &string_pos {
+            err.span_label(*pos, "string literal");
+
+        }
+        err.help("do not mix byte string literals and string literals");
+        err.multipart_suggestion(
+            "you can use byte string literals",
+            string_pos
+                .iter()
+                .map(|pos| (pos.shrink_to_lo(), "b".to_string()))
+                .collect(),
+        );
+        err.emit();
+    }
     let sp = sp.apply_mark(cx.current_expansion.mark);
-    base::MacEager::expr(cx.expr_str(sp, Symbol::intern(&accumulator)))
+    if b_accumulator.len() > 0 {
+        base::MacEager::expr(cx.expr_lit(
+            sp,
+            ast::LitKind::new_byte_str(b_accumulator),
+        ))
+    } else {
+        base::MacEager::expr(cx.expr_str(sp, Symbol::intern(&string_accumulator)))
+    }
 }

src/test/compile-fail/concat.rs

@@ -9,8 +9,8 @@
 // except according to those terms.
 
 fn main() {
-    concat!(b'f');  //~ ERROR: cannot concatenate a byte string literal
-    concat!(b"foo");  //~ ERROR: cannot concatenate a byte string literal
+    concat!(b'f');
+    concat!(b"foo");
     concat!(foo);   //~ ERROR: expected a literal
     concat!(foo()); //~ ERROR: expected a literal
 }

src/test/ui/byte-concat.rs

@@ -0,0 +1,15 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+fn main() {
+    let _ = concat!(b"abc", b"def");
+    let _ = concat!("abc", b"def", "ghi", b"jkl");
+    //~^ ERROR cannot concatenate a byte string literal with string literals
+}

src/test/ui/byte-concat.stderr

@@ -0,0 +1,18 @@
+error: cannot concatenate a byte string literal with string literals
+  --> $DIR/byte-concat.rs:13:28
+   |
+LL |     let _ = concat!("abc", b"def", "ghi", b"jkl");
+   |                     -----  ^^^^^^  -----  ^^^^^^ byte string literal
+   |                     |      |       |
+   |                     |      |       string literal
+   |                     |      byte string literal
+   |                     string literal
+   |
+   = help: do not mix byte string literals and string literals
+help: you can use byte string literals
+   |
+LL |     let _ = concat!(b"abc", b"def", b"ghi", b"jkl");
+   |                     ^               ^
+
+error: aborting due to previous error
+