Auto merge of #127313 - cjgillot:single-expect, r=jieyouxu

Rewrite lint_expectations in a single pass.

This PR aims at reducing the perf regression from #120924 (comment) with drive-by simplifications.

Basically, instead of using the lint level builder, which is slow, this PR splits `lint_expectations` logic in 2:
- listing the `LintExpectations` is done in `shallow_lint_levels_on`, on a per-owner basis;
- building the unstable->stable expectation id map is done by iterating on attributes.

r? ghost for perf

compiler/rustc_errors/src/diagnostic.rs

@@ -12,7 +12,7 @@ use rustc_lint_defs::{Applicability, LintExpectationId};
 use rustc_macros::{Decodable, Encodable};
 use rustc_span::source_map::Spanned;
 use rustc_span::symbol::Symbol;
-use rustc_span::{Span, DUMMY_SP};
+use rustc_span::{AttrId, Span, DUMMY_SP};
 use tracing::debug;
 
 use crate::snippet::Style;
@@ -356,24 +356,19 @@ impl DiagInner {
 
     pub(crate) fn update_unstable_expectation_id(
         &mut self,
-        unstable_to_stable: &FxIndexMap<LintExpectationId, LintExpectationId>,
+        unstable_to_stable: &FxIndexMap<AttrId, LintExpectationId>,
     ) {
         if let Level::Expect(expectation_id) | Level::ForceWarning(Some(expectation_id)) =
             &mut self.level
+            && let LintExpectationId::Unstable { attr_id, lint_index } = *expectation_id
         {
-            if expectation_id.is_stable() {
-                return;
-            }
-
             // The unstable to stable map only maps the unstable `AttrId` to a stable `HirId` with an attribute index.
             // The lint index inside the attribute is manually transferred here.
-            let lint_index = expectation_id.get_lint_index();
-            expectation_id.set_lint_index(None);
-            let mut stable_id = unstable_to_stable
-                .get(expectation_id)
-                .expect("each unstable `LintExpectationId` must have a matching stable id")
-                .normalize();
+            let Some(stable_id) = unstable_to_stable.get(&attr_id) else {
+                panic!("{expectation_id:?} must have a matching stable id")
+            };
 
+            let mut stable_id = *stable_id;
             stable_id.set_lint_index(lint_index);
             *expectation_id = stable_id;
         }

compiler/rustc_errors/src/lib.rs

@@ -69,7 +69,7 @@ use rustc_macros::{Decodable, Encodable};
 pub use rustc_span::fatal_error::{FatalError, FatalErrorMarker};
 use rustc_span::source_map::SourceMap;
 pub use rustc_span::ErrorGuaranteed;
-use rustc_span::{Loc, Span, DUMMY_SP};
+use rustc_span::{AttrId, Loc, Span, DUMMY_SP};
 pub use snippet::Style;
 // Used by external projects such as `rust-gpu`.
 // See https://github.com/rust-lang/rust/pull/115393.
@@ -1096,7 +1096,7 @@ impl<'a> DiagCtxtHandle<'a> {
 
     pub fn update_unstable_expectation_id(
         &self,
-        unstable_to_stable: &FxIndexMap<LintExpectationId, LintExpectationId>,
+        unstable_to_stable: FxIndexMap<AttrId, LintExpectationId>,
     ) {
         let mut inner = self.inner.borrow_mut();
         let diags = std::mem::take(&mut inner.unstable_expect_diagnostics);
@@ -1105,7 +1105,7 @@ impl<'a> DiagCtxtHandle<'a> {
         if !diags.is_empty() {
             inner.suppressed_expected_diag = true;
             for mut diag in diags.into_iter() {
-                diag.update_unstable_expectation_id(unstable_to_stable);
+                diag.update_unstable_expectation_id(&unstable_to_stable);
 
                 // Here the diagnostic is given back to `emit_diagnostic` where it was first
                 // intercepted. Now it should be processed as usual, since the unstable expectation
@@ -1117,11 +1117,11 @@ impl<'a> DiagCtxtHandle<'a> {
         inner
             .stashed_diagnostics
             .values_mut()
-            .for_each(|(diag, _guar)| diag.update_unstable_expectation_id(unstable_to_stable));
+            .for_each(|(diag, _guar)| diag.update_unstable_expectation_id(&unstable_to_stable));
         inner
             .future_breakage_diagnostics
             .iter_mut()
-            .for_each(|diag| diag.update_unstable_expectation_id(unstable_to_stable));
+            .for_each(|diag| diag.update_unstable_expectation_id(&unstable_to_stable));
     }
 
     /// This methods steals all [`LintExpectationId`]s that are stored inside
@@ -1567,7 +1567,7 @@ impl DiagCtxtInner {
                 if let LintExpectationId::Unstable { .. } = expect_id {
                     unreachable!(); // this case was handled at the top of this function
                 }
-                self.fulfilled_expectations.insert(expect_id.normalize());
+                self.fulfilled_expectations.insert(expect_id);
                 if let Expect(_) = diagnostic.level {
                     // Nothing emitted here for expected lints.
                     TRACK_DIAGNOSTIC(diagnostic, &mut |_| None);

compiler/rustc_lint/src/expect.rs

@@ -1,21 +1,66 @@
+use rustc_data_structures::fx::FxIndexMap;
+use rustc_hir::{HirId, CRATE_OWNER_ID};
+use rustc_middle::lint::LintExpectation;
 use rustc_middle::query::Providers;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::lint::builtin::UNFULFILLED_LINT_EXPECTATIONS;
-use rustc_session::lint::LintExpectationId;
+use rustc_session::lint::{Level, LintExpectationId};
 use rustc_span::Symbol;
 
 use crate::lints::{Expectation, ExpectationNote};
 
 pub(crate) fn provide(providers: &mut Providers) {
-    *providers = Providers { check_expectations, ..*providers };
+    *providers = Providers { lint_expectations, check_expectations, ..*providers };
+}
+
+fn lint_expectations(tcx: TyCtxt<'_>, (): ()) -> Vec<(LintExpectationId, LintExpectation)> {
+    let krate = tcx.hir_crate_items(());
+
+    let mut expectations = Vec::new();
+    let mut unstable_to_stable_ids = FxIndexMap::default();
+
+    let mut record_stable = |attr_id, hir_id, attr_index| {
+        let expect_id = LintExpectationId::Stable { hir_id, attr_index, lint_index: None };
+        unstable_to_stable_ids.entry(attr_id).or_insert(expect_id);
+    };
+    let mut push_expectations = |owner| {
+        let lints = tcx.shallow_lint_levels_on(owner);
+        if lints.expectations.is_empty() {
+            return;
+        }
+
+        expectations.extend_from_slice(&lints.expectations);
+
+        let attrs = tcx.hir_attrs(owner);
+        for &(local_id, attrs) in attrs.map.iter() {
+            // Some attributes appear multiple times in HIR, to ensure they are correctly taken
+            // into account where they matter. This means we cannot just associate the AttrId to
+            // the first HirId where we see it, but need to check it actually appears in a lint
+            // level.
+            // FIXME(cjgillot): Can this cause an attribute to appear in multiple expectation ids?
+            if !lints.specs.contains_key(&local_id) {
+                continue;
+            }
+            for (attr_index, attr) in attrs.iter().enumerate() {
+                let Some(Level::Expect(_)) = Level::from_attr(attr) else { continue };
+                record_stable(attr.id, HirId { owner, local_id }, attr_index.try_into().unwrap());
+            }
+        }
+    };
+
+    push_expectations(CRATE_OWNER_ID);
+    for owner in krate.owners() {
+        push_expectations(owner);
+    }
+
+    tcx.dcx().update_unstable_expectation_id(unstable_to_stable_ids);
+    expectations
 }
 
 fn check_expectations(tcx: TyCtxt<'_>, tool_filter: Option<Symbol>) {
     let lint_expectations = tcx.lint_expectations(());
     let fulfilled_expectations = tcx.dcx().steal_fulfilled_expectation_ids();
 
-    tracing::debug!(?lint_expectations, ?fulfilled_expectations);
-
     for (id, expectation) in lint_expectations {
         // This check will always be true, since `lint_expectations` only
         // holds stable ids