Auto merge of #113609 - nnethercote:maybe_lint_level_root_bounded-cache, r=cjgillot

Add a cache for `maybe_lint_level_root_bounded`

`maybe_lint_level_root_bounded` is called many times and traces node sub-paths many times. This PR adds a cache that lets many of these tracings be skipped, avoiding lots of calls to functions like `Map::attrs` and `Map::parent_id`.

r? `@cjgillot`

3 files changed, 71 insertions, 32 deletions

diff --git a/compiler/rustc_middle/src/lint.rs b/compiler/rustc_middle/src/lint.rs
index 81c1ae4f630..20230217afc 100644
--- a/compiler/rustc_middle/src/lint.rs
+++ b/compiler/rustc_middle/src/lint.rs
@@ -169,26 +169,6 @@ impl TyCtxt<'_> {
     pub fn lint_level_at_node(self, lint: &'static Lint, id: HirId) -> (Level, LintLevelSource) {
         self.shallow_lint_levels_on(id.owner).lint_level_id_at_node(self, LintId::of(lint), id)
     }
-
-    /// Walks upwards from `id` to find a node which might change lint levels with attributes.
-    /// It stops at `bound` and just returns it if reached.
-    pub fn maybe_lint_level_root_bounded(self, mut id: HirId, bound: HirId) -> HirId {
-        let hir = self.hir();
-        loop {
-            if id == bound {
-                return bound;
-            }
-
-            if hir.attrs(id).iter().any(|attr| Level::from_attr(attr).is_some()) {
-                return id;
-            }
-            let next = hir.parent_id(id);
-            if next == id {
-                bug!("lint traversal reached the root of the crate");
-            }
-            id = next;
-        }
-    }
 }
 
 /// This struct represents a lint expectation and holds all required information
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
index d828e71c7ac..a91ced4a211 100644
--- a/compiler/rustc_mir_build/src/build/mod.rs
+++ b/compiler/rustc_mir_build/src/build/mod.rs
@@ -10,6 +10,7 @@ use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::{GeneratorKind, Node};
+use rustc_index::bit_set::GrowableBitSet;
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_infer::infer::{InferCtxt, TyCtxtInferExt};
 use rustc_middle::hir::place::PlaceBase as HirPlaceBase;
@@ -215,6 +216,14 @@ struct Builder<'a, 'tcx> {
     unit_temp: Option<Place<'tcx>>,
 
     var_debug_info: Vec<VarDebugInfo<'tcx>>,
+
+    // A cache for `maybe_lint_level_roots_bounded`. That function is called
+    // repeatedly, and each time it effectively traces a path through a tree
+    // structure from a node towards the root, doing an attribute check on each
+    // node along the way. This cache records which nodes trace all the way to
+    // the root (most of them do) and saves us from retracing many sub-paths
+    // many times, and rechecking many nodes.
+    lint_level_roots_cache: GrowableBitSet<hir::ItemLocalId>,
 }
 
 type CaptureMap<'tcx> = SortedIndexMultiMap<usize, hir::HirId, Capture<'tcx>>;
@@ -725,6 +734,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             var_indices: Default::default(),
             unit_temp: None,
             var_debug_info: vec![],
+            lint_level_roots_cache: GrowableBitSet::new_empty(),
         };
 
         assert_eq!(builder.cfg.start_new_block(), START_BLOCK);
diff --git a/compiler/rustc_mir_build/src/build/scope.rs b/compiler/rustc_mir_build/src/build/scope.rs
index 72374102c8c..a96288a11e5 100644
--- a/compiler/rustc_mir_build/src/build/scope.rs
+++ b/compiler/rustc_mir_build/src/build/scope.rs
@@ -90,8 +90,8 @@ use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::middle::region;
 use rustc_middle::mir::*;
 use rustc_middle::thir::{Expr, LintLevel};
-
 use rustc_middle::ty::Ty;
+use rustc_session::lint::Level;
 use rustc_span::{Span, DUMMY_SP};
 
 #[derive(Debug)]
@@ -760,20 +760,25 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     ) {
         let (current_root, parent_root) =
             if self.tcx.sess.opts.unstable_opts.maximal_hir_to_mir_coverage {
-                // Some consumers of rustc need to map MIR locations back to HIR nodes. Currently the
-                // the only part of rustc that tracks MIR -> HIR is the `SourceScopeLocalData::lint_root`
-                // field that tracks lint levels for MIR locations. Normally the number of source scopes
-                // is limited to the set of nodes with lint annotations. The -Zmaximal-hir-to-mir-coverage
-                // flag changes this behavior to maximize the number of source scopes, increasing the
-                // granularity of the MIR->HIR mapping.
+                // Some consumers of rustc need to map MIR locations back to HIR nodes. Currently
+                // the the only part of rustc that tracks MIR -> HIR is the
+                // `SourceScopeLocalData::lint_root` field that tracks lint levels for MIR
+                // locations. Normally the number of source scopes is limited to the set of nodes
+                // with lint annotations. The -Zmaximal-hir-to-mir-coverage flag changes this
+                // behavior to maximize the number of source scopes, increasing the granularity of
+                // the MIR->HIR mapping.
                 (current_id, parent_id)
             } else {
-                // Use `maybe_lint_level_root_bounded` with `self.hir_id` as a bound
-                // to avoid adding Hir dependencies on our parents.
-                // We estimate the true lint roots here to avoid creating a lot of source scopes.
+                // Use `maybe_lint_level_root_bounded` to avoid adding Hir dependencies on our
+                // parents. We estimate the true lint roots here to avoid creating a lot of source
+                // scopes.
                 (
-                    self.tcx.maybe_lint_level_root_bounded(current_id, self.hir_id),
-                    self.tcx.maybe_lint_level_root_bounded(parent_id, self.hir_id),
+                    self.maybe_lint_level_root_bounded(current_id),
+                    if parent_id == self.hir_id {
+                        parent_id // this is very common
+                    } else {
+                        self.maybe_lint_level_root_bounded(parent_id)
+                    },
                 )
             };
 
@@ -783,6 +788,50 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
     }
 
+    /// Walks upwards from `orig_id` to find a node which might change lint levels with attributes.
+    /// It stops at `self.hir_id` and just returns it if reached.
+    fn maybe_lint_level_root_bounded(&mut self, orig_id: HirId) -> HirId {
+        // This assertion lets us just store `ItemLocalId` in the cache, rather
+        // than the full `HirId`.
+        assert_eq!(orig_id.owner, self.hir_id.owner);
+
+        let mut id = orig_id;
+        let hir = self.tcx.hir();
+        loop {
+            if id == self.hir_id {
+                // This is a moderately common case, mostly hit for previously unseen nodes.
+                break;
+            }
+
+            if hir.attrs(id).iter().any(|attr| Level::from_attr(attr).is_some()) {
+                // This is a rare case. It's for a node path that doesn't reach the root due to an
+                // intervening lint level attribute. This result doesn't get cached.
+                return id;
+            }
+
+            let next = hir.parent_id(id);
+            if next == id {
+                bug!("lint traversal reached the root of the crate");
+            }
+            id = next;
+
+            // This lookup is just an optimization; it can be removed without affecting
+            // functionality. It might seem strange to see this at the end of this loop, but the
+            // `orig_id` passed in to this function is almost always previously unseen, for which a
+            // lookup will be a miss. So we only do lookups for nodes up the parent chain, where
+            // cache lookups have a very high hit rate.
+            if self.lint_level_roots_cache.contains(id.local_id) {
+                break;
+            }
+        }
+
+        // `orig_id` traced to `self_id`; record this fact. If `orig_id` is a leaf node it will
+        // rarely (never?) subsequently be searched for, but it's hard to know if that is the case.
+        // The performance wins from the cache all come from caching non-leaf nodes.
+        self.lint_level_roots_cache.insert(orig_id.local_id);
+        self.hir_id
+    }
+
     /// Creates a new source scope, nested in the current one.
     pub(crate) fn new_source_scope(
         &mut self,