Preserve byte string literal raw-ness in AST

Author: petrochenkov

src/librustc/ich/impls_syntax.rs

@@ -166,7 +166,7 @@ impl_stable_hash_for_spanned!(::syntax::ast::LitKind);
 impl_stable_hash_for!(enum ::syntax::ast::LitKind {
     Str(value, style),
     Err(value),
-    ByteStr(value),
+    ByteStr(value, style),
     Byte(value),
     Char(value),
     Int(value, lit_int_type),

src/librustc_mir/hair/constant.rs

@@ -44,7 +44,7 @@ crate fn lit_to_const<'a, 'gcx, 'tcx>(
                 ty: tcx.types.err,
             });
         },
-        LitKind::ByteStr(ref data) => {
+        LitKind::ByteStr(ref data, _) => {
             let id = tcx.allocate_bytes(data);
             ConstValue::Scalar(Scalar::Ptr(id.into()))
         },

src/librustc_typeck/check/_match.rs

@@ -153,7 +153,7 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
                 // They can denote both statically and dynamically sized byte arrays
                 let mut pat_ty = ty;
                 if let hir::ExprKind::Lit(ref lt) = lt.node {
-                    if let ast::LitKind::ByteStr(_) = lt.node {
+                    if let ast::LitKind::ByteStr(..) = lt.node {
                         let expected_ty = self.structurally_resolved_type(pat.span, expected);
                         if let ty::Ref(_, r_ty, _) = expected_ty.sty {
                             if let ty::Slice(_) = r_ty.sty {

src/librustc_typeck/check/mod.rs

@@ -3092,7 +3092,7 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
 
         match lit.node {
             ast::LitKind::Str(..) => tcx.mk_static_str(),
-            ast::LitKind::ByteStr(ref v) => {
+            ast::LitKind::ByteStr(ref v, _) => {
                 tcx.mk_imm_ref(tcx.lifetimes.re_static,
                                tcx.mk_array(tcx.types.u8, v.len() as u64))
             }

src/libsyntax/ast.rs

@@ -1355,7 +1355,7 @@ pub enum LitKind {
     /// A string literal (`"foo"`).
     Str(Symbol, StrStyle),
     /// A byte string (`b"foo"`).
-    ByteStr(Lrc<Vec<u8>>),
+    ByteStr(Lrc<Vec<u8>>, StrStyle),
     /// A byte char (`b'f'`).
     Byte(u8),
     /// A character literal (`'a'`).
@@ -1384,7 +1384,7 @@ impl LitKind {
     /// Returns `true` if this literal is byte literal string.
     pub fn is_bytestr(&self) -> bool {
         match self {
-            LitKind::ByteStr(_) => true,
+            LitKind::ByteStr(..) => true,
             _ => false,
         }
     }

src/libsyntax/attr/mod.rs

@@ -635,11 +635,16 @@ impl LitKind {
             LitKind::Str(string, ast::StrStyle::Raw(n)) => {
                 Token::Literal(token::Lit::StrRaw(string, n), None)
             }
-            LitKind::ByteStr(ref bytes) => {
+            LitKind::ByteStr(ref bytes, ast::StrStyle::Cooked) => {
                 let string = bytes.iter().cloned().flat_map(ascii::escape_default)
                     .map(Into::<char>::into).collect::<String>();
                 Token::Literal(token::Lit::ByteStr(Symbol::intern(&string)), None)
             }
+            LitKind::ByteStr(ref bytes, ast::StrStyle::Raw(n)) => {
+                let string = String::from_utf8(bytes.to_vec())
+                    .expect("broken UTF-8 in a raw byte string literal");
+                Token::Literal(token::Lit::ByteStrRaw(Symbol::intern(&string), n), None)
+            }
             LitKind::Byte(byte) => {
                 let string: String = ascii::escape_default(byte).map(Into::<char>::into).collect();
                 Token::Literal(token::Lit::Byte(Symbol::intern(&string)), None)

src/libsyntax/ext/source_util.rs

@@ -166,7 +166,8 @@ pub fn expand_include_bytes(cx: &mut ExtCtxt<'_>, sp: Span, tts: &[tokenstream::
             };
             cx.source_map().new_source_file(file.into(), contents);
 
-            base::MacEager::expr(cx.expr_lit(sp, ast::LitKind::ByteStr(Lrc::new(bytes))))
+            let lit_kind = ast::LitKind::ByteStr(Lrc::new(bytes), ast::StrStyle::Raw(0));
+            base::MacEager::expr(cx.expr_lit(sp, lit_kind))
         },
         Err(e) => {
             cx.span_err(sp, &format!("couldn't read {}: {}", file.display(), e));

src/libsyntax/parse/mod.rs

@@ -419,10 +419,11 @@ crate fn lit_token(lit: token::Lit, suf: Option<Symbol>, diag: Option<(Span, &Ha
                 return (true, Some(LitKind::Err(i)));
             }
             buf.shrink_to_fit();
-            (true, Some(LitKind::ByteStr(Lrc::new(buf))))
+            (true, Some(LitKind::ByteStr(Lrc::new(buf), ast::StrStyle::Cooked)))
         }
-        token::ByteStrRaw(i, _) => {
-            (true, Some(LitKind::ByteStr(Lrc::new(i.to_string().into_bytes()))))
+        token::ByteStrRaw(i, n) => {
+            let bytes = Lrc::new(i.to_string().into_bytes());
+            (true, Some(LitKind::ByteStr(bytes, ast::StrStyle::Raw(n))))
         }
     }
 }

src/libsyntax/print/pprust.rs

@@ -611,6 +611,7 @@ pub trait PrintState<'a> {
         }
         match lit.node {
             ast::LitKind::Str(st, style) => self.print_string(&st.as_str(), style),
+            ast::LitKind::ByteStr(ref st, style) => self.print_byte_string(st, style),
             ast::LitKind::Err(st) => {
                 let st = st.as_str().escape_debug().to_string();
                 let mut res = String::with_capacity(st.len() + 2);
@@ -651,14 +652,6 @@ pub trait PrintState<'a> {
             ast::LitKind::Bool(val) => {
                 if val { self.writer().word("true") } else { self.writer().word("false") }
             }
-            ast::LitKind::ByteStr(ref v) => {
-                let mut escaped: String = String::new();
-                for &ch in v.iter() {
-                    escaped.extend(ascii::escape_default(ch)
-                                         .map(|c| c as char));
-                }
-                self.writer().word(format!("b\"{}\"", escaped))
-            }
         }
     }
 
@@ -677,6 +670,24 @@ pub trait PrintState<'a> {
         self.writer().word(st)
     }
 
+    fn print_byte_string(&mut self, st: &[u8], style: ast::StrStyle) -> io::Result<()> {
+        let st = match style {
+            ast::StrStyle::Cooked => {
+                let mut escaped = String::with_capacity(st.len());
+                for &c in st {
+                    escaped.extend(ascii::escape_default(c).map(|c| c as char));
+                }
+                format!("b\"{}\"", escaped)
+            }
+            ast::StrStyle::Raw(n) => {
+                let st = String::from_utf8(st.to_owned())
+                    .expect("broken UTF-8 in a raw byte string literal");
+                format!("br{delim}\"{string}\"{delim}", delim="#".repeat(n as usize), string=st)
+            }
+        };
+        self.writer().word(st)
+    }
+
     fn print_inner_attributes(&mut self,
                               attrs: &[ast::Attribute]) -> io::Result<()> {
         self.print_either_attributes(attrs, ast::AttrStyle::Inner, false, true)

src/test/pretty/raw-byte-str.rs

@@ -0,0 +1,4 @@
+// pp-exact
+
+fn main() { let x = br#"byt
+es"#; }