compiler/rustc_builtin_macros/src/format/expand.rs

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