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