Auto merge of rust-lang#103812 - clubby789:improve-include-bytes, r=petrochenkov

Delay `include_bytes` to AST lowering

Hopefully addresses rust-lang#65818.
This PR introduces a new `ExprKind::IncludedBytes` which stores the path and bytes of a file included with `include_bytes!()`. We can then create a literal from the bytes during AST lowering, which means we don't need to escape the bytes into valid UTF8 which is the cause of most of the overhead of embedding large binary blobs.

Author: bors

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;
             }