Introduce ExprKind::IncludedBytes

Author: clubby789

compiler/rustc_ast/src/ast.rs

@@ -1208,7 +1208,7 @@ impl Expr {
             ExprKind::Tup(_) => ExprPrecedence::Tup,
             ExprKind::Binary(op, ..) => ExprPrecedence::Binary(op.node),
             ExprKind::Unary(..) => ExprPrecedence::Unary,
-            ExprKind::Lit(_) => ExprPrecedence::Lit,
+            ExprKind::Lit(_) | ExprKind::IncludedBytes(..) => ExprPrecedence::Lit,
             ExprKind::Type(..) | ExprKind::Cast(..) => ExprPrecedence::Cast,
             ExprKind::Let(..) => ExprPrecedence::Let,
             ExprKind::If(..) => ExprPrecedence::If,
@@ -1446,6 +1446,12 @@ pub enum ExprKind {
     /// with an optional value to be returned.
     Yeet(Option<P<Expr>>),
 
+    /// Bytes included via `include_bytes!`
+    /// Added for optimization purposes to avoid the need to escape
+    /// large binary blobs - should always behave like [`ExprKind::Lit`]
+    /// with a `ByteStr` literal.
+    IncludedBytes(Lrc<[u8]>),
+
     /// Placeholder for an expression that wasn't syntactically well formed in some way.
     Err,
 }

compiler/rustc_ast/src/mut_visit.rs

@@ -1428,7 +1428,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
         }
         ExprKind::Try(expr) => vis.visit_expr(expr),
         ExprKind::TryBlock(body) => vis.visit_block(body),
-        ExprKind::Lit(_) | ExprKind::Err => {}
+        ExprKind::Lit(_) | ExprKind::IncludedBytes(..) | ExprKind::Err => {}
     }
     vis.visit_id(id);
     vis.visit_span(span);

compiler/rustc_ast/src/util/literal.rs

@@ -2,6 +2,7 @@
 
 use crate::ast::{self, Lit, LitKind};
 use crate::token::{self, Token};
+use rustc_data_structures::sync::Lrc;
 use rustc_lexer::unescape::{byte_from_char, unescape_byte, unescape_char, unescape_literal, Mode};
 use rustc_span::symbol::{kw, sym, Symbol};
 use rustc_span::Span;
@@ -231,6 +232,13 @@ impl Lit {
         Lit { token_lit: kind.to_token_lit(), kind, span }
     }
 
+    /// Recovers an AST literal from a string of bytes produced by `include_bytes!`.
+    /// This requires ASCII-escaping the string, which can result in poor performance
+    /// for very large strings of bytes.
+    pub fn from_included_bytes(bytes: &Lrc<[u8]>, span: Span) -> Lit {
+        Self::from_lit_kind(LitKind::ByteStr(bytes.clone()), span)
+    }
+
     /// Losslessly convert an AST literal into a token.
     pub fn to_token(&self) -> Token {
         let kind = match self.token_lit.kind {

compiler/rustc_ast/src/visit.rs

@@ -901,7 +901,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
         }
         ExprKind::Try(ref subexpression) => visitor.visit_expr(subexpression),
         ExprKind::TryBlock(ref body) => visitor.visit_block(body),
-        ExprKind::Lit(_) | ExprKind::Err => {}
+        ExprKind::Lit(_) | ExprKind::IncludedBytes(..) | ExprKind::Err => {}
     }
 
     visitor.visit_expr_post(expression)

compiler/rustc_ast_lowering/src/expr.rs

@@ -87,6 +87,10 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ExprKind::Lit(ref l) => {
                     hir::ExprKind::Lit(respan(self.lower_span(l.span), l.kind.clone()))
                 }
+                ExprKind::IncludedBytes(ref bytes) => hir::ExprKind::Lit(respan(
+                    self.lower_span(e.span),
+                    LitKind::ByteStr(bytes.clone()),
+                )),
                 ExprKind::Cast(ref expr, ref ty) => {
                     let expr = self.lower_expr(expr);
                     let ty =

compiler/rustc_ast_lowering/src/pat.rs

@@ -323,7 +323,10 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     // ```
     fn lower_expr_within_pat(&mut self, expr: &Expr, allow_paths: bool) -> &'hir hir::Expr<'hir> {
         match expr.kind {
-            ExprKind::Lit(..) | ExprKind::ConstBlock(..) | ExprKind::Err => {}
+            ExprKind::Lit(..)
+            | ExprKind::ConstBlock(..)
+            | ExprKind::IncludedBytes(..)
+            | ExprKind::Err => {}
             ExprKind::Path(..) if allow_paths => {}
             ExprKind::Unary(UnOp::Neg, ref inner) if matches!(inner.kind, ExprKind::Lit(_)) => {}
             _ => {

compiler/rustc_ast_pretty/src/pprust/state/expr.rs

@@ -322,6 +322,10 @@ impl<'a> State<'a> {
             ast::ExprKind::Lit(ref lit) => {
                 self.print_literal(lit);
             }
+            ast::ExprKind::IncludedBytes(ref bytes) => {
+                let lit = ast::Lit::from_included_bytes(bytes, expr.span);
+                self.print_literal(&lit)
+            }
             ast::ExprKind::Cast(ref expr, ref ty) => {
                 let prec = AssocOp::As.precedence() as i8;
                 self.print_expr_maybe_paren(expr, prec);

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -303,6 +303,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::Field(_, _)
             | ExprKind::ForLoop(_, _, _, _)
             | ExprKind::If(_, _, _)
+            | ExprKind::IncludedBytes(..)
             | ExprKind::InlineAsm(_)
             | ExprKind::Let(_, _, _)
             | ExprKind::Lit(_)

compiler/rustc_builtin_macros/src/concat.rs

@@ -43,6 +43,9 @@ pub fn expand_concat(
                     has_errors = true;
                 }
             },
+            ast::ExprKind::IncludedBytes(..) => {
+                cx.span_err(e.span, "cannot concatenate a byte string literal")
+            }
             ast::ExprKind::Err => {
                 has_errors = true;
             }

compiler/rustc_builtin_macros/src/concat_bytes.rs

@@ -108,6 +108,16 @@ fn handle_array_element(
                 None
             }
         },
+        ast::ExprKind::IncludedBytes(..) => {
+            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
+        }
         _ => {
             missing_literals.push(expr.span);
             None
@@ -167,6 +177,9 @@ pub fn expand_concat_bytes(
                     has_errors = true;
                 }
             },
+            ast::ExprKind::IncludedBytes(ref bytes) => {
+                accumulator.extend_from_slice(bytes);
+            }
             ast::ExprKind::Err => {
                 has_errors = true;
             }