Expand format_args!() in rust_ast_lowering.

6 files changed, 281 insertions, 1 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index e656fb3740b..5b1722dffb9 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -18,6 +18,7 @@
 //! - [`Attribute`]: Metadata associated with item.
 //! - [`UnOp`], [`BinOp`], and [`BinOpKind`]: Unary and binary operators.
 
+pub use crate::format::*;
 pub use crate::util::parser::ExprPrecedence;
 pub use GenericArgs::*;
 pub use UnsafeSource::*;
@@ -1269,6 +1270,7 @@ impl Expr {
             ExprKind::Try(..) => ExprPrecedence::Try,
             ExprKind::Yield(..) => ExprPrecedence::Yield,
             ExprKind::Yeet(..) => ExprPrecedence::Yeet,
+            ExprKind::FormatArgs(..) => ExprPrecedence::FormatArgs,
             ExprKind::Err => ExprPrecedence::Err,
         }
     }
@@ -1498,6 +1500,9 @@ pub enum ExprKind {
     /// with a `ByteStr` literal.
     IncludedBytes(Lrc<[u8]>),
 
+    /// A `format_args!()` expression.
+    FormatArgs(P<FormatArgs>),
+
     /// Placeholder for an expression that wasn't syntactically well formed in some way.
     Err,
 }
diff --git a/compiler/rustc_ast/src/format.rs b/compiler/rustc_ast/src/format.rs
new file mode 100644
index 00000000000..ce99c2b58b5
--- /dev/null
+++ b/compiler/rustc_ast/src/format.rs
@@ -0,0 +1,244 @@
+use crate::ptr::P;
+use crate::Expr;
+use rustc_data_structures::fx::FxHashMap;
+use rustc_span::symbol::{Ident, Symbol};
+use rustc_span::Span;
+
+// Definitions:
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+// └──────────────────────────────────────────────┘
+//                     FormatArgs
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//                                     └─────────┘
+//                                      argument
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//              └───────────────────┘
+//                     template
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//               └────┘└─────────┘└┘
+//                      pieces
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//               └────┘           └┘
+//                   literal pieces
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//                     └─────────┘
+//                     placeholder
+//
+// format_args!("hello {abc:.xyz$}!!", abc="world");
+//                      └─┘  └─┘
+//                      positions (could be names, numbers, empty, or `*`)
+
+/// (Parsed) format args.
+///
+/// Basically the "AST" for a complete `format_args!()`.
+///
+/// E.g., `format_args!("hello {name}");`.
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub struct FormatArgs {
+    pub span: Span,
+    pub template: Vec<FormatArgsPiece>,
+    pub arguments: FormatArguments,
+}
+
+/// A piece of a format template string.
+///
+/// E.g. "hello" or "{name}".
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub enum FormatArgsPiece {
+    Literal(Symbol),
+    Placeholder(FormatPlaceholder),
+}
+
+/// The arguments to format_args!().
+///
+/// E.g. `1, 2, name="ferris", n=3`,
+/// but also implicit captured arguments like `x` in `format_args!("{x}")`.
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub struct FormatArguments {
+    arguments: Vec<FormatArgument>,
+    num_unnamed_args: usize,
+    num_explicit_args: usize,
+    names: FxHashMap<Symbol, usize>,
+}
+
+impl FormatArguments {
+    pub fn new() -> Self {
+        Self {
+            arguments: Vec::new(),
+            names: FxHashMap::default(),
+            num_unnamed_args: 0,
+            num_explicit_args: 0,
+        }
+    }
+
+    pub fn add(&mut self, arg: FormatArgument) -> usize {
+        let index = self.arguments.len();
+        if let Some(name) = arg.kind.ident() {
+            self.names.insert(name.name, index);
+        } else if self.names.is_empty() {
+            // Only count the unnamed args before the first named arg.
+            // (Any later ones are errors.)
+            self.num_unnamed_args += 1;
+        }
+        if !matches!(arg.kind, FormatArgumentKind::Captured(..)) {
+            // This is an explicit argument.
+            // Make sure that all arguments so far are explcit.
+            assert_eq!(
+                self.num_explicit_args,
+                self.arguments.len(),
+                "captured arguments must be added last"
+            );
+            self.num_explicit_args += 1;
+        }
+        self.arguments.push(arg);
+        index
+    }
+
+    pub fn by_name(&self, name: Symbol) -> Option<(usize, &FormatArgument)> {
+        let i = *self.names.get(&name)?;
+        Some((i, &self.arguments[i]))
+    }
+
+    pub fn by_index(&self, i: usize) -> Option<&FormatArgument> {
+        (i < self.num_explicit_args).then(|| &self.arguments[i])
+    }
+
+    pub fn unnamed_args(&self) -> &[FormatArgument] {
+        &self.arguments[..self.num_unnamed_args]
+    }
+
+    pub fn named_args(&self) -> &[FormatArgument] {
+        &self.arguments[self.num_unnamed_args..self.num_explicit_args]
+    }
+
+    pub fn explicit_args(&self) -> &[FormatArgument] {
+        &self.arguments[..self.num_explicit_args]
+    }
+
+    pub fn all_args(&self) -> &[FormatArgument] {
+        &self.arguments[..]
+    }
+
+    pub fn all_args_mut(&mut self) -> &mut [FormatArgument] {
+        &mut self.arguments[..]
+    }
+}
+
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub struct FormatArgument {
+    pub kind: FormatArgumentKind,
+    pub expr: P<Expr>,
+}
+
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub enum FormatArgumentKind {
+    /// `format_args(…, arg)`
+    Normal,
+    /// `format_args(…, arg = 1)`
+    Named(Ident),
+    /// `format_args("… {arg} …")`
+    Captured(Ident),
+}
+
+impl FormatArgumentKind {
+    pub fn ident(&self) -> Option<Ident> {
+        match self {
+            &Self::Normal => None,
+            &Self::Named(id) => Some(id),
+            &Self::Captured(id) => Some(id),
+        }
+    }
+}
+
+#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub struct FormatPlaceholder {
+    /// Index into [`FormatArgs::arguments`].
+    pub argument: FormatArgPosition,
+    /// The span inside the format string for the full `{…}` placeholder.
+    pub span: Option<Span>,
+    /// `{}`, `{:?}`, or `{:x}`, etc.
+    pub format_trait: FormatTrait,
+    /// `{}` or `{:.5}` or `{:-^20}`, etc.
+    pub format_options: FormatOptions,
+}
+
+#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub struct FormatArgPosition {
+    /// Which argument this position refers to (Ok),
+    /// or would've referred to if it existed (Err).
+    pub index: Result<usize, usize>,
+    /// What kind of position this is. See [`FormatArgPositionKind`].
+    pub kind: FormatArgPositionKind,
+    /// The span of the name or number.
+    pub span: Option<Span>,
+}
+
+#[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub enum FormatArgPositionKind {
+    /// `{}` or `{:.*}`
+    Implicit,
+    /// `{1}` or `{:1$}` or `{:.1$}`
+    Number,
+    /// `{a}` or `{:a$}` or `{:.a$}`
+    Named,
+}
+
+#[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq, Hash)]
+pub enum FormatTrait {
+    /// `{}`
+    Display,
+    /// `{:?}`
+    Debug,
+    /// `{:e}`
+    LowerExp,
+    /// `{:E}`
+    UpperExp,
+    /// `{:o}`
+    Octal,
+    /// `{:p}`
+    Pointer,
+    /// `{:b}`
+    Binary,
+    /// `{:x}`
+    LowerHex,
+    /// `{:X}`
+    UpperHex,
+}
+
+#[derive(Clone, Encodable, Decodable, Default, Debug, PartialEq, Eq)]
+pub struct FormatOptions {
+    /// The width. E.g. `{:5}` or `{:width$}`.
+    pub width: Option<FormatCount>,
+    /// The precision. E.g. `{:.5}` or `{:.precision$}`.
+    pub precision: Option<FormatCount>,
+    /// The alignment. E.g. `{:>}` or `{:<}` or `{:^}`.
+    pub alignment: Option<FormatAlignment>,
+    /// The fill character. E.g. the `.` in `{:.>10}`.
+    pub fill: Option<char>,
+    /// The `+`, `-`, `0`, `#`, `x?` and `X?` flags.
+    pub flags: u32,
+}
+
+#[derive(Copy, Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub enum FormatAlignment {
+    /// `{:<}`
+    Left,
+    /// `{:>}`
+    Right,
+    /// `{:^}`
+    Center,
+}
+
+#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub enum FormatCount {
+    /// `{:5}` or `{:.5}`
+    Literal(usize),
+    /// `{:.*}`, `{:.5$}`, or `{:a$}`, etc.
+    Argument(FormatArgPosition),
+}
diff --git a/compiler/rustc_ast/src/lib.rs b/compiler/rustc_ast/src/lib.rs
index 9c1dfeb1a61..0f8ebcfdc15 100644
--- a/compiler/rustc_ast/src/lib.rs
+++ b/compiler/rustc_ast/src/lib.rs
@@ -42,6 +42,7 @@ pub mod ast_traits;
 pub mod attr;
 pub mod entry;
 pub mod expand;
+pub mod format;
 pub mod mut_visit;
 pub mod node_id;
 pub mod ptr;
@@ -51,6 +52,7 @@ pub mod visit;
 
 pub use self::ast::*;
 pub use self::ast_traits::{AstDeref, AstNodeWrapper, HasAttrs, HasNodeId, HasSpan, HasTokens};
+pub use self::format::*;
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index c572171e8f4..561274cc6f9 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -297,6 +297,10 @@ pub trait MutVisitor: Sized {
     fn visit_inline_asm_sym(&mut self, sym: &mut InlineAsmSym) {
         noop_visit_inline_asm_sym(sym, self)
     }
+
+    fn visit_format_args(&mut self, fmt: &mut FormatArgs) {
+        noop_visit_format_args(fmt, self)
+    }
 }
 
 /// Use a map-style function (`FnOnce(T) -> T`) to overwrite a `&mut T`. Useful
@@ -1284,6 +1288,15 @@ pub fn noop_visit_inline_asm_sym<T: MutVisitor>(
     vis.visit_path(path);
 }
 
+pub fn noop_visit_format_args<T: MutVisitor>(fmt: &mut FormatArgs, vis: &mut T) {
+    for arg in fmt.arguments.all_args_mut() {
+        if let FormatArgumentKind::Named(name) = &mut arg.kind {
+            vis.visit_ident(name);
+        }
+        vis.visit_expr(&mut arg.expr);
+    }
+}
+
 pub fn noop_visit_expr<T: MutVisitor>(
     Expr { kind, id, span, attrs, tokens }: &mut Expr,
     vis: &mut T,
@@ -1423,6 +1436,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             visit_opt(expr, |expr| vis.visit_expr(expr));
         }
         ExprKind::InlineAsm(asm) => vis.visit_inline_asm(asm),
+        ExprKind::FormatArgs(fmt) => vis.visit_format_args(fmt),
         ExprKind::MacCall(mac) => vis.visit_mac_call(mac),
         ExprKind::Struct(se) => {
             let StructExpr { qself, path, fields, rest } = se.deref_mut();
diff --git a/compiler/rustc_ast/src/util/parser.rs b/compiler/rustc_ast/src/util/parser.rs
index 819f1884a06..2db2ab5e811 100644
--- a/compiler/rustc_ast/src/util/parser.rs
+++ b/compiler/rustc_ast/src/util/parser.rs
@@ -271,6 +271,7 @@ pub enum ExprPrecedence {
     Try,
     InlineAsm,
     Mac,
+    FormatArgs,
 
     Array,
     Repeat,
@@ -335,7 +336,8 @@ impl ExprPrecedence {
             | ExprPrecedence::Index
             | ExprPrecedence::Try
             | ExprPrecedence::InlineAsm
-            | ExprPrecedence::Mac => PREC_POSTFIX,
+            | ExprPrecedence::Mac
+            | ExprPrecedence::FormatArgs => PREC_POSTFIX,
 
             // Never need parens
             ExprPrecedence::Array
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index df7145a722a..cb5c17084ec 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -242,6 +242,9 @@ pub trait Visitor<'ast>: Sized {
     fn visit_inline_asm(&mut self, asm: &'ast InlineAsm) {
         walk_inline_asm(self, asm)
     }
+    fn visit_format_args(&mut self, fmt: &'ast FormatArgs) {
+        walk_format_args(self, fmt)
+    }
     fn visit_inline_asm_sym(&mut self, sym: &'ast InlineAsmSym) {
         walk_inline_asm_sym(self, sym)
     }
@@ -756,6 +759,15 @@ pub fn walk_inline_asm_sym<'a, V: Visitor<'a>>(visitor: &mut V, sym: &'a InlineA
     visitor.visit_path(&sym.path, sym.id);
 }
 
+pub fn walk_format_args<'a, V: Visitor<'a>>(visitor: &mut V, fmt: &'a FormatArgs) {
+    for arg in fmt.arguments.all_args() {
+        if let FormatArgumentKind::Named(name) = arg.kind {
+            visitor.visit_ident(name);
+        }
+        visitor.visit_expr(&arg.expr);
+    }
+}
+
 pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
     walk_list!(visitor, visit_attribute, expression.attrs.iter());
 
@@ -895,6 +907,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
         ExprKind::MacCall(mac) => visitor.visit_mac_call(mac),
         ExprKind::Paren(subexpression) => visitor.visit_expr(subexpression),
         ExprKind::InlineAsm(asm) => visitor.visit_inline_asm(asm),
+        ExprKind::FormatArgs(f) => visitor.visit_format_args(f),
         ExprKind::Yield(optional_expression) => {
             walk_list!(visitor, visit_expr, optional_expression);
         }