Rollup merge of #154745 - chenyukang:yukang-fix-span-api, r=nnethercote

Replace span_look_ahead with span_followed_by

While reviewing that PR https://github.com/rust-lang/rust/pull/154703#discussion_r3031067780, I found that magic number 100, let's remove it, and seems `span_followed_by` is a better name.

compiler/rustc_mir_build/src/thir/pattern/check_match.rs

@@ -1100,8 +1100,7 @@ fn report_arm_reachability<'p, 'tcx>(
             let arm_span = cx.tcx.hir_span(hir_id);
             let whole_arm_span = if is_match_arm {
                 // If the arm is followed by a comma, extend the span to include it.
-                let with_whitespace = sm.span_extend_while_whitespace(arm_span);
-                if let Some(comma) = sm.span_look_ahead(with_whitespace, ",", Some(1)) {
+                if let Some(comma) = sm.span_followed_by(arm_span, ",") {
                     Some(arm_span.to(comma))
                 } else {
                     Some(arm_span)

compiler/rustc_resolve/src/ident.rs

@@ -1467,11 +1467,9 @@ fn validate_res_from_ribs(
                                         // `const name: Ty = expr;`. This is a heuristic, it will
                                         // break down in the presence of macros.
                                         let sm = self.tcx.sess.source_map();
-                                        let type_span = match sm.span_look_ahead(
-                                            original_rib_ident_def.span,
-                                            ":",
-                                            None,
-                                        ) {
+                                        let type_span = match sm
+                                            .span_followed_by(original_rib_ident_def.span, ":")
+                                        {
                                             None => {
                                                 Some(original_rib_ident_def.span.shrink_to_hi())
                                             }

compiler/rustc_resolve/src/late/diagnostics.rs

@@ -1989,10 +1989,25 @@ fn followed_by_brace(&self, span: Span) -> (bool, Option<Span>) {
         // where a brace being opened means a block is being started. Look
         // ahead for the next text to see if `span` is followed by a `{`.
         let sm = self.r.tcx.sess.source_map();
-        if let Some(followed_brace_span) = sm.span_look_ahead(span, "{", Some(50)) {
+        if let Some(open_brace_span) = sm.span_followed_by(span, "{") {
             // In case this could be a struct literal that needs to be surrounded
             // by parentheses, find the appropriate span.
-            let close_brace_span = sm.span_look_ahead(followed_brace_span, "}", Some(50));
+            let close_brace_span =
+                sm.span_to_next_source(open_brace_span).ok().and_then(|next_source| {
+                    // Find the matching `}` accounting for nested braces.
+                    let mut depth: u32 = 1;
+                    let offset = next_source.char_indices().find_map(|(i, c)| {
+                        match c {
+                            '{' => depth += 1,
+                            '}' if depth == 1 => return Some(i),
+                            '}' => depth -= 1,
+                            _ => {}
+                        }
+                        None
+                    })?;
+                    let start = open_brace_span.hi() + rustc_span::BytePos(offset as u32);
+                    Some(open_brace_span.with_lo(start).with_hi(start + rustc_span::BytePos(1)))
+                });
             let closing_brace = close_brace_span.map(|sp| span.to(sp));
             (true, closing_brace)
         } else {
@@ -4110,8 +4125,7 @@ fn add_missing_lifetime_specifiers_label<'a>(
                 let sugg: String = std::iter::repeat_n(existing_name.as_str(), lt.count)
                     .intersperse(", ")
                     .collect();
-                let is_empty_brackets =
-                    source_map.span_look_ahead(lt.span, ">", Some(50)).is_some();
+                let is_empty_brackets = source_map.span_followed_by(lt.span, ">").is_some();
                 let sugg = if is_empty_brackets { sugg } else { format!("{sugg}, ") };
                 (lt.span.shrink_to_hi(), sugg)
             }

compiler/rustc_span/src/source_map.rs

@@ -964,21 +964,13 @@ pub fn next_point(&self, sp: Span) -> Span {
         Span::new(BytePos(start_of_next_point), end_of_next_point, sp.ctxt, None)
     }
 
-    /// Check whether span is followed by some specified expected string in limit scope
-    pub fn span_look_ahead(&self, span: Span, expect: &str, limit: Option<usize>) -> Option<Span> {
-        let mut sp = span;
-        for _ in 0..limit.unwrap_or(100_usize) {
-            sp = self.next_point(sp);
-            if let Ok(ref snippet) = self.span_to_snippet(sp) {
-                if snippet == expect {
-                    return Some(sp);
-                }
-                if snippet.chars().any(|c| !c.is_whitespace()) {
-                    break;
-                }
-            }
-        }
-        None
+    /// Check whether span is followed by some specified target string, ignoring whitespace.
+    /// *Only suitable for diagnostics.*
+    pub fn span_followed_by(&self, span: Span, target: &str) -> Option<Span> {
+        let span = self.span_extend_while_whitespace(span);
+        self.span_to_next_source(span).ok()?.strip_prefix(target).map(|_| {
+            Span::new(span.hi(), span.hi() + BytePos(target.len() as u32), span.ctxt(), None)
+        })
     }
 
     /// Finds the width of the character, either before or after the end of provided span,

compiler/rustc_span/src/source_map/tests.rs

@@ -752,6 +752,25 @@ fn test_next_point() {
     assert!(sm.span_to_snippet(span).is_err());
 }
 
+#[test]
+fn test_span_followed_by_stops_at_end_of_file() {
+    let sm = SourceMap::new(FilePathMapping::empty());
+    sm.new_source_file(filename(&sm, "example.rs"), "x".to_string());
+
+    let span = Span::with_root_ctxt(BytePos(0), BytePos(1));
+    assert_eq!(sm.span_followed_by(span, "y"), None);
+}
+
+#[test]
+fn test_span_followed_by_skips_whitespace() {
+    let sm = SourceMap::new(FilePathMapping::empty());
+    sm.new_source_file(filename(&sm, "example.rs"), "x \n yz".to_string());
+
+    let span = Span::with_root_ctxt(BytePos(0), BytePos(1));
+    let span = sm.span_followed_by(span, "yz").unwrap();
+    assert_eq!(sm.span_to_snippet(span), Ok("yz".to_string()));
+}
+
 #[cfg(target_os = "linux")]
 #[test]
 fn read_binary_file_handles_lying_stat() {

compiler/rustc_trait_selection/src/error_reporting/traits/fulfillment_errors.rs

@@ -459,7 +459,7 @@ pub fn report_selection_error(
                                 (cand.self_ty().kind(), main_trait_predicate.self_ty().skip_binder().kind())
                             {
                                 // Wrap method receivers and `&`-references in parens
-                                let suggestion = if self.tcx.sess.source_map().span_look_ahead(span, ".", Some(50)).is_some() {
+                                let suggestion = if self.tcx.sess.source_map().span_followed_by(span, ".").is_some() {
                                     vec![
                                         (span.shrink_to_lo(), format!("(")),
                                         (span.shrink_to_hi(), format!(" as {})", cand.self_ty())),

tests/ui/error-codes/E0423-struct-literal-comment.rs

@@ -0,0 +1,18 @@
+#[derive(PartialEq)]
+struct T { pub x: i32 }
+
+#[derive(PartialEq)]
+struct U { }
+
+fn main() {
+    // Parser will report an error here
+    if T { x: 10 } == T {} {}
+    //~^ ERROR struct literals are not allowed here
+    //~| ERROR expected value, found struct `T`
+
+    // Regression test for the `followed_by_brace` helper:
+    // comments inside the braces should not suppress the parenthesized struct literal suggestion.
+    if U { /* keep comment here */ } == U {}
+    //~^ ERROR E0423
+    //~| ERROR expected expression, found `==`
+}

tests/ui/error-codes/E0423-struct-literal-comment.stderr

@@ -0,0 +1,42 @@
+error: struct literals are not allowed here
+  --> $DIR/E0423-struct-literal-comment.rs:9:8
+   |
+LL |     if T { x: 10 } == T {} {}
+   |        ^^^^^^^^^^^
+   |
+help: surround the struct literal with parentheses
+   |
+LL |     if (T { x: 10 }) == T {} {}
+   |        +           +
+
+error: expected expression, found `==`
+  --> $DIR/E0423-struct-literal-comment.rs:15:38
+   |
+LL |     if U { /* keep comment here */ } == U {}
+   |                                      ^^ expected expression
+
+error[E0423]: expected value, found struct `T`
+  --> $DIR/E0423-struct-literal-comment.rs:9:23
+   |
+LL |     if T { x: 10 } == T {} {}
+   |                       ^ not a value
+   |
+help: surround the struct literal with parentheses
+   |
+LL |     if T { x: 10 } == (T {}) {}
+   |                       +    +
+
+error[E0423]: expected value, found struct `U`
+  --> $DIR/E0423-struct-literal-comment.rs:15:8
+   |
+LL |     if U { /* keep comment here */ } == U {}
+   |        ^ not a value
+   |
+help: surround the struct literal with parentheses
+   |
+LL |     if (U { /* keep comment here */ }) == U {}
+   |        +                             +
+
+error: aborting due to 4 previous errors
+
+For more information about this error, try `rustc --explain E0423`.