compiler/rustc_ast_lowering/src/format.rs

   1 use super::LoweringContext;
   2 use rustc_ast as ast;
   3 use rustc_ast::visit::{self, Visitor};
   4 use rustc_ast::*;
   5 use rustc_data_structures::fx::FxIndexSet;
   6 use rustc_hir as hir;
   7 use rustc_span::{
   8     sym,
   9     symbol::{kw, Ident},
  10     Span,
  11 };
  12
  13 impl<'hir> LoweringContext<'_, 'hir> {
  14     pub(crate) fn lower_format_args(&mut self, sp: Span, fmt: &FormatArgs) -> hir::ExprKind<'hir> {
  15         expand_format_args(self, sp, fmt)
  16     }
  17 }
  18
  19 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
  20 enum ArgumentType {
  21     Format(FormatTrait),
  22     Usize,
  23 }
  24
  25 /// Generate a hir expression representing an argument to a format_args invocation.
  26 ///
  27 /// Generates:
  28 ///
  29 /// ```text
  30 ///     <core::fmt::ArgumentV1>::new_…(arg)
  31 /// ```
  32 fn make_argument<'hir>(
  33     ctx: &mut LoweringContext<'_, 'hir>,
  34     sp: Span,
  35     arg: &'hir hir::Expr<'hir>,
  36     ty: ArgumentType,
  37 ) -> hir::Expr<'hir> {
  38     use ArgumentType::*;
  39     use FormatTrait::*;
  40     let new_fn = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
  41         sp,
  42         hir::LangItem::FormatArgument,
  43         match ty {
  44             Format(Display) => sym::new_display,
  45             Format(Debug) => sym::new_debug,
  46             Format(LowerExp) => sym::new_lower_exp,
  47             Format(UpperExp) => sym::new_upper_exp,
  48             Format(Octal) => sym::new_octal,
  49             Format(Pointer) => sym::new_pointer,
  50             Format(Binary) => sym::new_binary,
  51             Format(LowerHex) => sym::new_lower_hex,
  52             Format(UpperHex) => sym::new_upper_hex,
  53             Usize => sym::from_usize,
  54         },
  55     ));
  56     ctx.expr_call_mut(sp, new_fn, std::slice::from_ref(arg))
  57 }
  58
  59 /// Generate a hir expression for a format_args Count.
  60 ///
  61 /// Generates:
  62 ///
  63 /// ```text
  64 ///     <core::fmt::rt::v1::Count>::Is(…)
  65 /// ```
  66 ///
  67 /// or
  68 ///
  69 /// ```text
  70 ///     <core::fmt::rt::v1::Count>::Param(…)
  71 /// ```
  72 ///
  73 /// or
  74 ///
  75 /// ```text
  76 ///     <core::fmt::rt::v1::Count>::Implied
  77 /// ```
  78 fn make_count<'hir>(
  79     ctx: &mut LoweringContext<'_, 'hir>,
  80     sp: Span,
  81     count: &Option<FormatCount>,
  82     argmap: &mut FxIndexSet<(usize, ArgumentType)>,
  83 ) -> hir::Expr<'hir> {
  84     match count {
  85         Some(FormatCount::Literal(n)) => {
  86             let count_is = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
  87                 sp,
  88                 hir::LangItem::FormatCount,
  89                 sym::Is,
  90             ));
  91             let value = ctx.arena.alloc_from_iter([ctx.expr_usize(sp, *n)]);
  92             ctx.expr_call_mut(sp, count_is, value)
  93         }
  94         Some(FormatCount::Argument(arg)) => {
  95             if let Ok(arg_index) = arg.index {
  96                 let (i, _) = argmap.insert_full((arg_index, ArgumentType::Usize));
  97                 let count_param = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
  98                     sp,
  99                     hir::LangItem::FormatCount,
 100                     sym::Param,
 101                 ));
 102                 let value = ctx.arena.alloc_from_iter([ctx.expr_usize(sp, i)]);
 103                 ctx.expr_call_mut(sp, count_param, value)
 104             } else {
 105                 ctx.expr(sp, hir::ExprKind::Err)
 106             }
 107         }
 108         None => ctx.expr_lang_item_type_relative(sp, hir::LangItem::FormatCount, sym::Implied),
 109     }
 110 }
 111
 112 /// Generate a hir expression for a format_args placeholder specification.
 113 ///
 114 /// Generates
 115 ///
 116 /// ```text
 117 ///     <core::fmt::rt::v1::Argument::new(
 118 ///         …usize, // position
 119 ///         '…', // fill
 120 ///         <core::fmt::rt::v1::Alignment>::…, // alignment
 121 ///         …u32, // flags
 122 ///         <core::fmt::rt::v1::Count::…>, // width
 123 ///         <core::fmt::rt::v1::Count::…>, // precision
 124 ///     )
 125 /// ```
 126 fn make_format_spec<'hir>(
 127     ctx: &mut LoweringContext<'_, 'hir>,
 128     sp: Span,
 129     placeholder: &FormatPlaceholder,
 130     argmap: &mut FxIndexSet<(usize, ArgumentType)>,
 131 ) -> hir::Expr<'hir> {
 132     let position = match placeholder.argument.index {
 133         Ok(arg_index) => {
 134             let (i, _) =
 135                 argmap.insert_full((arg_index, ArgumentType::Format(placeholder.format_trait)));
 136             ctx.expr_usize(sp, i)
 137         }
 138         Err(_) => ctx.expr(sp, hir::ExprKind::Err),
 139     };
 140     let fill = ctx.expr_char(sp, placeholder.format_options.fill.unwrap_or(' '));
 141     let align = ctx.expr_lang_item_type_relative(
 142         sp,
 143         hir::LangItem::FormatAlignment,
 144         match placeholder.format_options.alignment {
 145             Some(FormatAlignment::Left) => sym::Left,
 146             Some(FormatAlignment::Right) => sym::Right,
 147             Some(FormatAlignment::Center) => sym::Center,
 148             None => sym::Unknown,
 149         },
 150     );
 151     let flags = ctx.expr_u32(sp, placeholder.format_options.flags);
 152     let prec = make_count(ctx, sp, &placeholder.format_options.precision, argmap);
 153     let width = make_count(ctx, sp, &placeholder.format_options.width, argmap);
 154     let format_placeholder_new = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
 155         sp,
 156         hir::LangItem::FormatPlaceholder,
 157         sym::new,
 158     ));
 159     let args = ctx.arena.alloc_from_iter([position, fill, align, flags, prec, width]);
 160     ctx.expr_call_mut(sp, format_placeholder_new, args)
 161 }
 162
 163 fn expand_format_args<'hir>(
 164     ctx: &mut LoweringContext<'_, 'hir>,
 165     macsp: Span,
 166     fmt: &FormatArgs,
 167 ) -> hir::ExprKind<'hir> {
 168     let lit_pieces =
 169         ctx.arena.alloc_from_iter(fmt.template.iter().enumerate().filter_map(|(i, piece)| {
 170             match piece {
 171                 &FormatArgsPiece::Literal(s) => Some(ctx.expr_str(fmt.span, s)),
 172                 &FormatArgsPiece::Placeholder(_) => {
 173                     // Inject empty string before placeholders when not already preceded by a literal piece.
 174                     if i == 0 || matches!(fmt.template[i - 1], FormatArgsPiece::Placeholder(_)) {
 175                         Some(ctx.expr_str(fmt.span, kw::Empty))
 176                     } else {
 177                         None
 178                     }
 179                 }
 180             }
 181         }));
 182     let lit_pieces = ctx.expr_array_ref(fmt.span, lit_pieces);
 183
 184     // Whether we'll use the `Arguments::new_v1_formatted` form (true),
 185     // or the `Arguments::new_v1` form (false).
 186     let mut use_format_options = false;
 187
 188     // Create a list of all _unique_ (argument, format trait) combinations.
 189     // E.g. "{0} {0:x} {0} {1}" -> [(0, Display), (0, LowerHex), (1, Display)]
 190     let mut argmap = FxIndexSet::default();
 191     for piece in &fmt.template {
 192         let FormatArgsPiece::Placeholder(placeholder) = piece else { continue };
 193         if placeholder.format_options != Default::default() {
 194             // Can't use basic form if there's any formatting options.
 195             use_format_options = true;
 196         }
 197         if let Ok(index) = placeholder.argument.index {
 198             if !argmap.insert((index, ArgumentType::Format(placeholder.format_trait))) {
 199                 // Duplicate (argument, format trait) combination,
 200                 // which we'll only put once in the args array.
 201                 use_format_options = true;
 202             }
 203         }
 204     }
 205
 206     let format_options = use_format_options.then(|| {
 207         // Generate:
 208         //     &[format_spec_0, format_spec_1, format_spec_2]
 209         let elements: Vec<_> = fmt
 210             .template
 211             .iter()
 212             .filter_map(|piece| {
 213                 let FormatArgsPiece::Placeholder(placeholder) = piece else { return None };
 214                 Some(make_format_spec(ctx, macsp, placeholder, &mut argmap))
 215             })
 216             .collect();
 217         ctx.expr_array_ref(macsp, ctx.arena.alloc_from_iter(elements))
 218     });
 219
 220     let arguments = fmt.arguments.all_args();
 221
 222     // If the args array contains exactly all the original arguments once,
 223     // in order, we can use a simple array instead of a `match` construction.
 224     // However, if there's a yield point in any argument except the first one,
 225     // we don't do this, because an ArgumentV1 cannot be kept across yield points.
 226     //
 227     // This is an optimization, speeding up compilation about 1-2% in some cases.
 228     // See https://github.com/rust-lang/rust/pull/106770#issuecomment-1380790609
 229     let use_simple_array = argmap.len() == arguments.len()
 230         && argmap.iter().enumerate().all(|(i, &(j, _))| i == j)
 231         && arguments.iter().skip(1).all(|arg| !may_contain_yield_point(&arg.expr));
 232
 233     let args = if use_simple_array {
 234         // Generate:
 235         //     &[
 236         //         <core::fmt::ArgumentV1>::new_display(&arg0),
 237         //         <core::fmt::ArgumentV1>::new_lower_hex(&arg1),
 238         //         <core::fmt::ArgumentV1>::new_debug(&arg2),
 239         //         …
 240         //     ]
 241         let elements: Vec<_> = arguments
 242             .iter()
 243             .zip(argmap)
 244             .map(|(arg, (_, ty))| {
 245                 let sp = arg.expr.span.with_ctxt(macsp.ctxt());
 246                 let arg = ctx.lower_expr(&arg.expr);
 247                 let ref_arg = ctx.arena.alloc(ctx.expr(
 248                     sp,
 249                     hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg),
 250                 ));
 251                 make_argument(ctx, sp, ref_arg, ty)
 252             })
 253             .collect();
 254         ctx.expr_array_ref(macsp, ctx.arena.alloc_from_iter(elements))
 255     } else {
 256         // Generate:
 257         //     &match (&arg0, &arg1, &…) {
 258         //         args => [
 259         //             <core::fmt::ArgumentV1>::new_display(args.0),
 260         //             <core::fmt::ArgumentV1>::new_lower_hex(args.1),
 261         //             <core::fmt::ArgumentV1>::new_debug(args.0),
 262         //             …
 263         //         ]
 264         //     }
 265         let args_ident = Ident::new(sym::args, macsp);
 266         let (args_pat, args_hir_id) = ctx.pat_ident(macsp, args_ident);
 267         let args = ctx.arena.alloc_from_iter(argmap.iter().map(|&(arg_index, ty)| {
 268             if let Some(arg) = arguments.get(arg_index) {
 269                 let sp = arg.expr.span.with_ctxt(macsp.ctxt());
 270                 let args_ident_expr = ctx.expr_ident(macsp, args_ident, args_hir_id);
 271                 let arg = ctx.arena.alloc(ctx.expr(
 272                     sp,
 273                     hir::ExprKind::Field(
 274                         args_ident_expr,
 275                         Ident::new(sym::integer(arg_index), macsp),
 276                     ),
 277                 ));
 278                 make_argument(ctx, sp, arg, ty)
 279             } else {
 280                 ctx.expr(macsp, hir::ExprKind::Err)
 281             }
 282         }));
 283         let elements: Vec<_> = arguments
 284             .iter()
 285             .map(|arg| {
 286                 let arg_expr = ctx.lower_expr(&arg.expr);
 287                 ctx.expr(
 288                     arg.expr.span.with_ctxt(macsp.ctxt()),
 289                     hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg_expr),
 290                 )
 291             })
 292             .collect();
 293         let args_tuple = ctx
 294             .arena
 295             .alloc(ctx.expr(macsp, hir::ExprKind::Tup(ctx.arena.alloc_from_iter(elements))));
 296         let array = ctx.arena.alloc(ctx.expr(macsp, hir::ExprKind::Array(args)));
 297         let match_arms = ctx.arena.alloc_from_iter([ctx.arm(args_pat, array)]);
 298         let match_expr = ctx.arena.alloc(ctx.expr_match(
 299             macsp,
 300             args_tuple,
 301             match_arms,
 302             hir::MatchSource::FormatArgs,
 303         ));
 304         ctx.expr(
 305             macsp,
 306             hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, match_expr),
 307         )
 308     };
 309
 310     if let Some(format_options) = format_options {
 311         // Generate:
 312         //     <core::fmt::Arguments>::new_v1_formatted(
 313         //         lit_pieces,
 314         //         args,
 315         //         format_options,
 316         //         unsafe { ::core::fmt::UnsafeArg::new() }
 317         //     )
 318         let new_v1_formatted = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
 319             macsp,
 320             hir::LangItem::FormatArguments,
 321             sym::new_v1_formatted,
 322         ));
 323         let unsafe_arg_new = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
 324             macsp,
 325             hir::LangItem::FormatUnsafeArg,
 326             sym::new,
 327         ));
 328         let unsafe_arg_new_call = ctx.expr_call(macsp, unsafe_arg_new, &[]);
 329         let hir_id = ctx.next_id();
 330         let unsafe_arg = ctx.expr_block(ctx.arena.alloc(hir::Block {
 331             stmts: &[],
 332             expr: Some(unsafe_arg_new_call),
 333             hir_id,
 334             rules: hir::BlockCheckMode::UnsafeBlock(hir::UnsafeSource::CompilerGenerated),
 335             span: macsp,
 336             targeted_by_break: false,
 337         }));
 338         let args = ctx.arena.alloc_from_iter([lit_pieces, args, format_options, unsafe_arg]);
 339         hir::ExprKind::Call(new_v1_formatted, args)
 340     } else {
 341         // Generate:
 342         //     <core::fmt::Arguments>::new_v1(
 343         //         lit_pieces,
 344         //         args,
 345         //     )
 346         let new_v1 = ctx.arena.alloc(ctx.expr_lang_item_type_relative(
 347             macsp,
 348             hir::LangItem::FormatArguments,
 349             sym::new_v1,
 350         ));
 351         let new_args = ctx.arena.alloc_from_iter([lit_pieces, args]);
 352         hir::ExprKind::Call(new_v1, new_args)
 353     }
 354 }
 355
 356 fn may_contain_yield_point(e: &ast::Expr) -> bool {
 357     struct MayContainYieldPoint(bool);
 358
 359     impl Visitor<'_> for MayContainYieldPoint {
 360         fn visit_expr(&mut self, e: &ast::Expr) {
 361             if let ast::ExprKind::Await(_) | ast::ExprKind::Yield(_) = e.kind {
 362                 self.0 = true;
 363             } else {
 364                 visit::walk_expr(self, e);
 365             }
 366         }
 367
 368         fn visit_mac_call(&mut self, _: &ast::MacCall) {
 369             // Macros should be expanded at this point.
 370             unreachable!("unexpanded macro in ast lowering");
 371         }
 372
 373         fn visit_item(&mut self, _: &ast::Item) {
 374             // Do not recurse into nested items.
 375         }
 376     }
 377
 378     let mut visitor = MayContainYieldPoint(false);
 379     visitor.visit_expr(e);
 380     visitor.0
 381 }