Fix rustc_parse_format precision & width spans

3 files changed, 84 insertions, 65 deletions

diff --git a/compiler/rustc_builtin_macros/src/format.rs b/compiler/rustc_builtin_macros/src/format.rs
index 08026c9d357..fe3d8e6ce98 100644
--- a/compiler/rustc_builtin_macros/src/format.rs
+++ b/compiler/rustc_builtin_macros/src/format.rs
@@ -413,7 +413,7 @@ impl<'a, 'b> Context<'a, 'b> {
     /// Verifies one piece of a parse string, and remembers it if valid.
     /// All errors are not emitted as fatal so we can continue giving errors
     /// about this and possibly other format strings.
-    fn verify_piece(&mut self, p: &parse::Piece<'_>) {
+    fn verify_piece(&mut self, p: &parse::Piece<'a>) {
         match *p {
             parse::String(..) => {}
             parse::NextArgument(ref arg) => {
@@ -433,6 +433,11 @@ impl<'a, 'b> Context<'a, 'b> {
                 let has_precision = arg.format.precision != Count::CountImplied;
                 let has_width = arg.format.width != Count::CountImplied;
 
+                if has_precision || has_width {
+                    // push before named params are resolved to aid diagnostics
+                    self.arg_with_formatting.push(arg.format);
+                }
+
                 // argument second, if it's an implicit positional parameter
                 // it's written second, so it should come after width/precision.
                 let pos = match arg.position {
@@ -581,7 +586,11 @@ impl<'a, 'b> Context<'a, 'b> {
         let mut zero_based_note = false;
 
         let count = self.pieces.len()
-            + self.arg_with_formatting.iter().filter(|fmt| fmt.precision_span.is_some()).count();
+            + self
+                .arg_with_formatting
+                .iter()
+                .filter(|fmt| matches!(fmt.precision, parse::CountIsParam(_)))
+                .count();
         if self.names.is_empty() && !numbered_position_args && count != self.num_args() {
             e = self.ecx.struct_span_err(
                 sp,
@@ -647,7 +656,7 @@ impl<'a, 'b> Context<'a, 'b> {
                             + self
                                 .arg_with_formatting
                                 .iter()
-                                .filter(|fmt| fmt.precision_span.is_some())
+                                .filter(|fmt| matches!(fmt.precision, parse::CountIsParam(_)))
                                 .count();
                         e.span_label(
                             span,
@@ -899,26 +908,22 @@ impl<'a, 'b> Context<'a, 'b> {
                     },
                     position_span: arg.position_span,
                     format: parse::FormatSpec {
-                        fill: arg.format.fill,
+                        fill: None,
                         align: parse::AlignUnknown,
                         flags: 0,
                         precision: parse::CountImplied,
-                        precision_span: None,
+                        precision_span: arg.format.precision_span,
                         width: parse::CountImplied,
-                        width_span: None,
+                        width_span: arg.format.width_span,
                         ty: arg.format.ty,
                         ty_span: arg.format.ty_span,
                     },
                 };
 
                 let fill = arg.format.fill.unwrap_or(' ');
-
                 let pos_simple = arg.position.index() == simple_arg.position.index();
 
-                if arg.format.precision_span.is_some() || arg.format.width_span.is_some() {
-                    self.arg_with_formatting.push(arg.format);
-                }
-                if !pos_simple || arg.format != simple_arg.format || fill != ' ' {
+                if !pos_simple || arg.format != simple_arg.format {
                     self.all_pieces_simple = false;
                 }
 
diff --git a/compiler/rustc_parse_format/src/lib.rs b/compiler/rustc_parse_format/src/lib.rs
index e4842d2afb7..b63a173cc29 100644
--- a/compiler/rustc_parse_format/src/lib.rs
+++ b/compiler/rustc_parse_format/src/lib.rs
@@ -264,9 +264,7 @@ impl<'a> Iterator for Parser<'a> {
             }
         } else {
             if self.is_literal {
-                let start = self.to_span_index(self.cur_line_start);
-                let end = self.to_span_index(self.input.len());
-                let span = start.to(end);
+                let span = self.span(self.cur_line_start, self.input.len());
                 if self.line_spans.last() != Some(&span) {
                     self.line_spans.push(span);
                 }
@@ -384,6 +382,12 @@ impl<'a> Parser<'a> {
         InnerOffset(raw + pos + 1)
     }
 
+    fn span(&self, start_pos: usize, end_pos: usize) -> InnerSpan {
+        let start = self.to_span_index(start_pos);
+        let end = self.to_span_index(end_pos);
+        start.to(end)
+    }
+
     /// Forces consumption of the specified character. If the character is not
     /// found, an error is emitted.
     fn must_consume(&mut self, c: char) -> Option<usize> {
@@ -472,9 +476,7 @@ impl<'a> Parser<'a> {
                     return &self.input[start..pos];
                 }
                 '\n' if self.is_literal => {
-                    let start = self.to_span_index(self.cur_line_start);
-                    let end = self.to_span_index(pos);
-                    self.line_spans.push(start.to(end));
+                    self.line_spans.push(self.span(self.cur_line_start, pos));
                     self.cur_line_start = pos + 1;
                     self.cur.next();
                 }
@@ -537,6 +539,10 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn current_pos(&mut self) -> usize {
+        if let Some(&(pos, _)) = self.cur.peek() { pos } else { self.input.len() }
+    }
+
     /// Parses a format specifier at the current position, returning all of the
     /// relevant information in the `FormatSpec` struct.
     fn format(&mut self) -> FormatSpec<'a> {
@@ -590,39 +596,37 @@ impl<'a> Parser<'a> {
             // no '0' flag and '0$' as the width instead.
             if let Some(end) = self.consume_pos('$') {
                 spec.width = CountIsParam(0);
-
-                if let Some((pos, _)) = self.cur.peek().cloned() {
-                    spec.width_span = Some(self.to_span_index(pos - 2).to(self.to_span_index(pos)));
-                }
+                spec.width_span = Some(self.span(end - 1, end + 1));
                 havewidth = true;
-                spec.width_span = Some(self.to_span_index(end - 1).to(self.to_span_index(end + 1)));
             } else {
                 spec.flags |= 1 << (FlagSignAwareZeroPad as u32);
             }
         }
+
         if !havewidth {
-            let width_span_start = if let Some((pos, _)) = self.cur.peek() { *pos } else { 0 };
-            let (w, sp) = self.count(width_span_start);
-            spec.width = w;
-            spec.width_span = sp;
+            let start = self.current_pos();
+            spec.width = self.count(start);
+            if spec.width != CountImplied {
+                let end = self.current_pos();
+                spec.width_span = Some(self.span(start, end));
+            }
         }
 
         if let Some(start) = self.consume_pos('.') {
-            if let Some(end) = self.consume_pos('*') {
+            if self.consume('*') {
                 // Resolve `CountIsNextParam`.
                 // We can do this immediately as `position` is resolved later.
                 let i = self.curarg;
                 self.curarg += 1;
                 spec.precision = CountIsParam(i);
-                spec.precision_span =
-                    Some(self.to_span_index(start).to(self.to_span_index(end + 1)));
             } else {
-                let (p, sp) = self.count(start);
-                spec.precision = p;
-                spec.precision_span = sp;
+                spec.precision = self.count(start + 1);
             }
+            let end = self.current_pos();
+            spec.precision_span = Some(self.span(start, end));
         }
-        let ty_span_start = self.cur.peek().map(|(pos, _)| *pos);
+
+        let ty_span_start = self.current_pos();
         // Optional radix followed by the actual format specifier
         if self.consume('x') {
             if self.consume('?') {
@@ -642,11 +646,9 @@ impl<'a> Parser<'a> {
             spec.ty = "?";
         } else {
             spec.ty = self.word();
-            let ty_span_end = self.cur.peek().map(|(pos, _)| *pos);
             if !spec.ty.is_empty() {
-                spec.ty_span = ty_span_start
-                    .and_then(|s| ty_span_end.map(|e| (s, e)))
-                    .map(|(start, end)| self.to_span_index(start).to(self.to_span_index(end)));
+                let ty_span_end = self.current_pos();
+                spec.ty_span = Some(self.span(ty_span_start, ty_span_end));
             }
         }
         spec
@@ -670,13 +672,11 @@ impl<'a> Parser<'a> {
             return spec;
         }
 
-        let ty_span_start = self.cur.peek().map(|(pos, _)| *pos);
+        let ty_span_start = self.current_pos();
         spec.ty = self.word();
-        let ty_span_end = self.cur.peek().map(|(pos, _)| *pos);
         if !spec.ty.is_empty() {
-            spec.ty_span = ty_span_start
-                .and_then(|s| ty_span_end.map(|e| (s, e)))
-                .map(|(start, end)| self.to_span_index(start).to(self.to_span_index(end)));
+            let ty_span_end = self.current_pos();
+            spec.ty_span = Some(self.span(ty_span_start, ty_span_end));
         }
 
         spec
@@ -685,26 +685,21 @@ impl<'a> Parser<'a> {
     /// Parses a `Count` parameter at the current position. This does not check
     /// for 'CountIsNextParam' because that is only used in precision, not
     /// width.
-    fn count(&mut self, start: usize) -> (Count<'a>, Option<InnerSpan>) {
+    fn count(&mut self, start: usize) -> Count<'a> {
         if let Some(i) = self.integer() {
-            if let Some(end) = self.consume_pos('$') {
-                let span = self.to_span_index(start).to(self.to_span_index(end + 1));
-                (CountIsParam(i), Some(span))
-            } else {
-                (CountIs(i), None)
-            }
+            if self.consume('$') { CountIsParam(i) } else { CountIs(i) }
         } else {
             let tmp = self.cur.clone();
             let word = self.word();
             if word.is_empty() {
                 self.cur = tmp;
-                (CountImplied, None)
+                CountImplied
             } else if let Some(end) = self.consume_pos('$') {
-                let span = self.to_span_index(start + 1).to(self.to_span_index(end));
-                (CountIsName(word, span), None)
+                let name_span = self.span(start, end);
+                CountIsName(word, name_span)
             } else {
                 self.cur = tmp;
-                (CountImplied, None)
+                CountImplied
             }
         }
     }
@@ -737,7 +732,7 @@ impl<'a> Parser<'a> {
                 "invalid argument name `_`",
                 "invalid argument name",
                 "argument name cannot be a single underscore",
-                self.to_span_index(start).to(self.to_span_index(end)),
+                self.span(start, end),
             );
         }
         word
diff --git a/compiler/rustc_parse_format/src/tests.rs b/compiler/rustc_parse_format/src/tests.rs
index 57853069610..44ef0cd0eb5 100644
--- a/compiler/rustc_parse_format/src/tests.rs
+++ b/compiler/rustc_parse_format/src/tests.rs
@@ -1,5 +1,6 @@
 use super::*;
 
+#[track_caller]
 fn same(fmt: &'static str, p: &[Piece<'static>]) {
     let parser = Parser::new(fmt, None, None, false, ParseMode::Format);
     assert_eq!(parser.collect::<Vec<Piece<'static>>>(), p);
@@ -190,9 +191,9 @@ fn format_counts() {
                 align: AlignUnknown,
                 flags: 0,
                 precision: CountImplied,
-                width: CountIs(10),
                 precision_span: None,
-                width_span: None,
+                width: CountIs(10),
+                width_span: Some(InnerSpan { start: 3, end: 5 }),
                 ty: "x",
                 ty_span: None,
             },
@@ -208,9 +209,9 @@ fn format_counts() {
                 align: AlignUnknown,
                 flags: 0,
                 precision: CountIs(10),
+                precision_span: Some(InnerSpan { start: 6, end: 9 }),
                 width: CountIsParam(10),
-                precision_span: None,
-                width_span: Some(InnerSpan::new(3, 6)),
+                width_span: Some(InnerSpan { start: 3, end: 6 }),
                 ty: "x",
                 ty_span: None,
             },
@@ -226,9 +227,9 @@ fn format_counts() {
                 align: AlignUnknown,
                 flags: 0,
                 precision: CountIs(10),
+                precision_span: Some(InnerSpan { start: 6, end: 9 }),
                 width: CountIsParam(0),
-                precision_span: None,
-                width_span: Some(InnerSpan::new(4, 6)),
+                width_span: Some(InnerSpan { start: 4, end: 6 }),
                 ty: "x",
                 ty_span: None,
             },
@@ -244,8 +245,8 @@ fn format_counts() {
                 align: AlignUnknown,
                 flags: 0,
                 precision: CountIsParam(0),
+                precision_span: Some(InnerSpan { start: 3, end: 5 }),
                 width: CountImplied,
-                precision_span: Some(InnerSpan::new(3, 5)),
                 width_span: None,
                 ty: "x",
                 ty_span: None,
@@ -279,15 +280,33 @@ fn format_counts() {
                 fill: None,
                 align: AlignUnknown,
                 flags: 0,
-                precision: CountIsName("b", InnerSpan::new(6, 7)),
-                width: CountIsName("a", InnerSpan::new(4, 4)),
-                precision_span: None,
-                width_span: None,
+                precision: CountIsName("b", InnerSpan { start: 6, end: 7 }),
+                precision_span: Some(InnerSpan { start: 5, end: 8 }),
+                width: CountIsName("a", InnerSpan { start: 3, end: 4 }),
+                width_span: Some(InnerSpan { start: 3, end: 5 }),
                 ty: "?",
                 ty_span: None,
             },
         })],
     );
+    same(
+        "{:.4}",
+        &[NextArgument(Argument {
+            position: ArgumentImplicitlyIs(0),
+            position_span: InnerSpan { start: 2, end: 2 },
+            format: FormatSpec {
+                fill: None,
+                align: AlignUnknown,
+                flags: 0,
+                precision: CountIs(4),
+                precision_span: Some(InnerSpan { start: 3, end: 5 }),
+                width: CountImplied,
+                width_span: None,
+                ty: "",
+                ty_span: None,
+            },
+        })],
+    )
 }
 #[test]
 fn format_flags() {