diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs index db0dd2fcc6191..cf171f032f39a 100644 --- a/compiler/rustc_ast_lowering/src/expr.rs +++ b/compiler/rustc_ast_lowering/src/expr.rs @@ -2261,6 +2261,16 @@ impl<'hir> LoweringContext<'_, 'hir> { self.expr(span, hir::ExprKind::Lit(Spanned { node: LitKind::Bool(val), span })) } + pub(super) fn expr_usize_literal(&mut self, span: Span, val: u128) -> hir::Expr<'hir> { + self.expr( + span, + hir::ExprKind::Lit(Spanned { + node: LitKind::Int(val.into(), LitIntType::Unsigned(UintTy::Usize)), + span, + }), + ) + } + pub(super) fn expr(&mut self, span: Span, kind: hir::ExprKind<'hir>) -> hir::Expr<'hir> { let hir_id = self.next_id(); hir::Expr { hir_id, kind, span: self.lower_span(span) } diff --git a/compiler/rustc_ast_lowering/src/format.rs b/compiler/rustc_ast_lowering/src/format.rs index b9974edea71a4..daa05fc3c5c1b 100644 --- a/compiler/rustc_ast_lowering/src/format.rs +++ b/compiler/rustc_ast_lowering/src/format.rs @@ -308,9 +308,13 @@ fn expand_format_args<'hir>( // See library/core/src/fmt/mod.rs for the format string encoding format. + let mut starts_with_placeholder = false; + let mut total_literal_len = 0; for (i, piece) in template.iter().enumerate() { match piece { &FormatArgsPiece::Literal(sym) => { + total_literal_len += sym.as_str().len(); + // Coalesce adjacent literal pieces. if let Some(FormatArgsPiece::Literal(_)) = template.get(i + 1) { incomplete_lit.push_str(sym.as_str()); @@ -357,6 +361,10 @@ fn expand_format_args<'hir>( incomplete_lit.clear(); } FormatArgsPiece::Placeholder(p) => { + if total_literal_len == 0 { + starts_with_placeholder = true; + } + // Push the start byte and remember its index so we can set the option bits later. let i = bytecode.len(); bytecode.push(0xC0); @@ -494,9 +502,26 @@ fn expand_format_args<'hir>( ) }; + // `Arguments::estimated_capacity()` is used by `alloc::fmt::format` to reduce + // the number of reallocations. This is a somewhat reasonable heuristic given + // the limited information available at this time, but there is definitely + // room for improvement here. + let estimated_capacity = if starts_with_placeholder && total_literal_len < 16 { + // If the format string starts with a placeholder, + // don't preallocate anything, unless length + // of literal pieces is significant. + 0 + } else { + // There are some placeholders, so any additional push + // will reallocate the string. To avoid that, + // we're "pre-doubling" the capacity here. + (total_literal_len as u128).wrapping_mul(2) + }; + let estimated_capacity = ctx.expr_usize_literal(macsp, estimated_capacity); + // Generate: // unsafe { - // ::new(b"…", &args) + // ::new(b"…", &args, estimated_capacity) // } let template = ctx.expr_byte_str(macsp, ByteSymbol::intern(&bytecode)); let call = { @@ -506,7 +531,7 @@ fn expand_format_args<'hir>( sym::new, )); let args = ctx.expr_ref(macsp, args); - let new_args = ctx.arena.alloc_from_iter([template, args]); + let new_args = ctx.arena.alloc_from_iter([template, args, estimated_capacity]); ctx.expr_call(macsp, new, new_args) }; let call = hir::ExprKind::Block( diff --git a/library/core/src/fmt/mod.rs b/library/core/src/fmt/mod.rs index 7a3d874f3b6f4..b75f3f6ad8e60 100644 --- a/library/core/src/fmt/mod.rs +++ b/library/core/src/fmt/mod.rs @@ -716,6 +716,7 @@ impl<'a> Formatter<'a> { pub struct Arguments<'a> { template: NonNull, args: NonNull>, + estimated_capacity: usize, } /// Used by the format_args!() macro to create a fmt::Arguments object. @@ -729,9 +730,16 @@ impl<'a> Arguments<'a> { pub unsafe fn new( template: &'a [u8; N], args: &'a [rt::Argument<'a>; M], + estimated_capacity: usize, ) -> Arguments<'a> { // SAFETY: Responsibility of the caller. - unsafe { Arguments { template: mem::transmute(template), args: mem::transmute(args) } } + unsafe { + Arguments { + template: mem::transmute(template), + args: mem::transmute(args), + estimated_capacity, + } + } } // Same as `from_str`, but not const. @@ -752,57 +760,7 @@ impl<'a> Arguments<'a> { /// when using `format!`. Note: this is neither the lower nor upper bound. #[inline] pub fn estimated_capacity(&self) -> usize { - if let Some(s) = self.as_str() { - return s.len(); - } - // Iterate over the template, counting the length of literal pieces. - let mut length = 0usize; - let mut starts_with_placeholder = false; - let mut template = self.template; - loop { - // SAFETY: We can assume the template is valid. - unsafe { - let n = template.read(); - template = template.add(1); - if n == 0 { - // End of template. - break; - } else if n < 128 { - // Short literal string piece. - length += n as usize; - template = template.add(n as usize); - } else if n == 128 { - // Long literal string piece. - let len = usize::from(u16::from_le_bytes(template.cast_array().read())); - length += len; - template = template.add(2 + len); - } else { - assert_unchecked(n >= 0xC0); - // Placeholder piece. - if length == 0 { - starts_with_placeholder = true; - } - // Skip remainder of placeholder: - let skip = (n & 1 != 0) as usize * 4 // flags (32 bit) - + (n & 2 != 0) as usize * 2 // width (16 bit) - + (n & 4 != 0) as usize * 2 // precision (16 bit) - + (n & 8 != 0) as usize * 2; // arg_index (16 bit) - template = template.add(skip as usize); - } - } - } - - if starts_with_placeholder && length < 16 { - // If the format string starts with a placeholder, - // don't preallocate anything, unless length - // of literal pieces is significant. - 0 - } else { - // There are some placeholders, so any additional push - // will reallocate the string. To avoid that, - // we're "pre-doubling" the capacity here. - length.wrapping_mul(2) - } + self.estimated_capacity } } @@ -818,6 +776,7 @@ impl<'a> Arguments<'a> { Arguments { template: mem::transmute(s.as_ptr()), args: mem::transmute(s.len() << 1 | 1), + estimated_capacity: s.len(), } } }