Simplify bound var resolution.

cjgillot · cjgillot · commit 5dd65e484c1d · 2024-08-18T16:10:54.000Z
diff --git a/compiler/rustc_hir_analysis/src/collect/resolve_bound_vars.rs b/compiler/rustc_hir_analysis/src/collect/resolve_bound_vars.rs
@@ -11,11 +11,14 @@ use std::fmt;
 
 use rustc_ast::visit::walk_list;
 use rustc_data_structures::fx::{FxHashSet, FxIndexMap, FxIndexSet};
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
 use rustc_hir::def_id::LocalDefId;
 use rustc_hir::intravisit::{self, Visitor};
-use rustc_hir::{GenericArg, GenericParam, GenericParamKind, HirId, HirIdMap, LifetimeName, Node};
+use rustc_hir::{
+    GenericArg, GenericParam, GenericParamKind, HirId, ItemLocalMap, LifetimeName, Node,
+};
 use rustc_macros::extension;
 use rustc_middle::hir::nested_filter;
 use rustc_middle::middle::resolve_bound_vars::*;
@@ -74,15 +77,15 @@ impl ResolvedArg {
 struct NamedVarMap {
     // maps from every use of a named (not anonymous) bound var to a
     // `ResolvedArg` describing how that variable is bound
-    defs: HirIdMap<ResolvedArg>,
+    defs: ItemLocalMap<ResolvedArg>,
 
     // Maps relevant hir items to the bound vars on them. These include:
     // - function defs
     // - function pointers
     // - closures
     // - trait refs
     // - bound types (like `T` in `for<'a> T<'a>: Foo`)
-    late_bound_vars: HirIdMap<Vec<ty::BoundVariableKind>>,
+    late_bound_vars: ItemLocalMap<Vec<ty::BoundVariableKind>>,
 }
 
 struct BoundVarContext<'a, 'tcx> {
@@ -225,10 +228,10 @@ pub(crate) fn provide(providers: &mut Providers) {
     *providers = Providers {
         resolve_bound_vars,
 
-        named_variable_map: |tcx, id| tcx.resolve_bound_vars(id).defs.get(&id),
+        named_variable_map: |tcx, id| &tcx.resolve_bound_vars(id).defs,
         is_late_bound_map,
         object_lifetime_default,
-        late_bound_vars_map: |tcx, id| tcx.resolve_bound_vars(id).late_bound_vars.get(&id),
+        late_bound_vars_map: |tcx, id| &tcx.resolve_bound_vars(id).late_bound_vars,
 
         ..*providers
     };
@@ -265,16 +268,12 @@ fn resolve_bound_vars(tcx: TyCtxt<'_>, local_def_id: hir::OwnerId) -> ResolveBou
         hir::OwnerNode::Synthetic => unreachable!(),
     }
 
-    let mut rl = ResolveBoundVars::default();
-
-    for (hir_id, v) in named_variable_map.defs {
-        let map = rl.defs.entry(hir_id.owner).or_default();
-        map.insert(hir_id.local_id, v);
-    }
-    for (hir_id, v) in named_variable_map.late_bound_vars {
-        let map = rl.late_bound_vars.entry(hir_id.owner).or_default();
-        map.insert(hir_id.local_id, v);
-    }
+    let defs = named_variable_map.defs.into_sorted_stable_ord();
+    let late_bound_vars = named_variable_map.late_bound_vars.into_sorted_stable_ord();
+    let rl = ResolveBoundVars {
+        defs: SortedMap::from_presorted_elements(defs),
+        late_bound_vars: SortedMap::from_presorted_elements(late_bound_vars),
+    };
 
     debug!(?rl.defs);
     debug!(?rl.late_bound_vars);
@@ -339,7 +338,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
                 Scope::Binder { hir_id, .. } => {
                     // Nested poly trait refs have the binders concatenated
                     let mut full_binders =
-                        self.map.late_bound_vars.entry(*hir_id).or_default().clone();
+                        self.map.late_bound_vars.entry(hir_id.local_id).or_default().clone();
                     full_binders.extend(supertrait_bound_vars);
                     break (full_binders, BinderScopeType::Concatenating);
                 }
@@ -678,7 +677,7 @@ impl<'a, 'tcx> Visitor<'tcx> for BoundVarContext<'a, 'tcx> {
             hir::TyKind::Ref(lifetime_ref, ref mt) => {
                 self.visit_lifetime(lifetime_ref);
                 let scope = Scope::ObjectLifetimeDefault {
-                    lifetime: self.map.defs.get(&lifetime_ref.hir_id).cloned(),
+                    lifetime: self.map.defs.get(&lifetime_ref.hir_id.local_id).cloned(),
                     s: self.scope,
                 };
                 self.with(scope, |this| this.visit_ty(mt.ty));
@@ -705,7 +704,7 @@ impl<'a, 'tcx> Visitor<'tcx> for BoundVarContext<'a, 'tcx> {
                     // and ban them. Type variables instantiated inside binders aren't
                     // well-supported at the moment, so this doesn't work.
                     // In the future, this should be fixed and this error should be removed.
-                    let def = self.map.defs.get(&lifetime.hir_id).copied();
+                    let def = self.map.defs.get(&lifetime.hir_id.local_id).copied();
                     let Some(ResolvedArg::LateBound(_, _, lifetime_def_id)) = def else { continue };
                     let Some(lifetime_def_id) = lifetime_def_id.as_local() else { continue };
                     let lifetime_hir_id = self.tcx.local_def_id_to_hir_id(lifetime_def_id);
@@ -843,7 +842,7 @@ impl<'a, 'tcx> Visitor<'tcx> for BoundVarContext<'a, 'tcx> {
             let bound_vars: Vec<_> =
                 self.tcx.fn_sig(sig_id).skip_binder().bound_vars().iter().collect();
             let hir_id = self.tcx.local_def_id_to_hir_id(def_id);
-            self.map.late_bound_vars.insert(hir_id, bound_vars);
+            self.map.late_bound_vars.insert(hir_id.local_id, bound_vars);
         }
         self.visit_fn_like_elision(fd.inputs, output, matches!(fk, intravisit::FnKind::Closure));
         intravisit::walk_fn_kind(self, fk);
@@ -1016,10 +1015,10 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
     }
 
     fn record_late_bound_vars(&mut self, hir_id: HirId, binder: Vec<ty::BoundVariableKind>) {
-        if let Some(old) = self.map.late_bound_vars.insert(hir_id, binder) {
+        if let Some(old) = self.map.late_bound_vars.insert(hir_id.local_id, binder) {
             bug!(
                 "overwrote bound vars for {hir_id:?}:\nold={old:?}\nnew={:?}",
-                self.map.late_bound_vars[&hir_id]
+                self.map.late_bound_vars[&hir_id.local_id]
             )
         }
     }
@@ -1394,9 +1393,9 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
                         kind.descr(param_def_id.to_def_id())
                     ),
                 };
-                self.map.defs.insert(hir_id, ResolvedArg::Error(guar));
+                self.map.defs.insert(hir_id.local_id, ResolvedArg::Error(guar));
             } else {
-                self.map.defs.insert(hir_id, def);
+                self.map.defs.insert(hir_id.local_id, def);
             }
             return;
         }
@@ -1429,7 +1428,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
                             bug!("unexpected def-kind: {}", kind.descr(param_def_id.to_def_id()))
                         }
                     });
-                    self.map.defs.insert(hir_id, ResolvedArg::Error(guar));
+                    self.map.defs.insert(hir_id.local_id, ResolvedArg::Error(guar));
                     return;
                 }
                 Scope::Root { .. } => break,
@@ -1539,7 +1538,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
                     // This index can be used with `generic_args` since `parent_count == 0`.
                     let index = generics.param_def_id_to_index[&param_def_id] as usize;
                     generic_args.args.get(index).and_then(|arg| match arg {
-                        GenericArg::Lifetime(lt) => map.defs.get(&lt.hir_id).copied(),
+                        GenericArg::Lifetime(lt) => map.defs.get(&lt.hir_id.local_id).copied(),
                         _ => None,
                     })
                 }
@@ -1829,7 +1828,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
     #[instrument(level = "debug", skip(self))]
     fn insert_lifetime(&mut self, lifetime_ref: &'tcx hir::Lifetime, def: ResolvedArg) {
         debug!(span = ?lifetime_ref.ident.span);
-        self.map.defs.insert(lifetime_ref.hir_id, def);
+        self.map.defs.insert(lifetime_ref.hir_id.local_id, def);
     }
 
     /// Sometimes we resolve a lifetime, but later find that it is an
@@ -1840,8 +1839,7 @@ impl<'a, 'tcx> BoundVarContext<'a, 'tcx> {
         lifetime_ref: &'tcx hir::Lifetime,
         bad_def: ResolvedArg,
     ) {
-        // FIXME(#120456) - is `swap_remove` correct?
-        let old_value = self.map.defs.swap_remove(&lifetime_ref.hir_id);
+        let old_value = self.map.defs.remove(&lifetime_ref.hir_id.local_id);
         assert_eq!(old_value, Some(bad_def));
     }
 }
diff --git a/compiler/rustc_middle/src/middle/resolve_bound_vars.rs b/compiler/rustc_middle/src/middle/resolve_bound_vars.rs
@@ -1,9 +1,9 @@
 //! Name resolution for lifetimes and late-bound type and const variables: type declarations.
 
-use rustc_data_structures::fx::FxIndexMap;
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir::def_id::DefId;
-use rustc_hir::{ItemLocalId, OwnerId};
+use rustc_hir::ItemLocalId;
 use rustc_macros::{Decodable, Encodable, HashStable, TyDecodable, TyEncodable};
 
 use crate::ty;
@@ -47,11 +47,11 @@ pub enum ObjectLifetimeDefault {
 
 /// Maps the id of each lifetime reference to the lifetime decl
 /// that it corresponds to.
-#[derive(Default, HashStable, Debug)]
+#[derive(HashStable, Debug)]
 pub struct ResolveBoundVars {
     /// Maps from every use of a named (not anonymous) lifetime to a
     /// `Region` describing how that region is bound
-    pub defs: FxIndexMap<OwnerId, FxIndexMap<ItemLocalId, ResolvedArg>>,
+    pub defs: SortedMap<ItemLocalId, ResolvedArg>,
 
-    pub late_bound_vars: FxIndexMap<OwnerId, FxIndexMap<ItemLocalId, Vec<ty::BoundVariableKind>>>,
+    pub late_bound_vars: SortedMap<ItemLocalId, Vec<ty::BoundVariableKind>>,
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -16,6 +16,7 @@ use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_ast::expand::StrippedCfgItem;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::svh::Svh;
 use rustc_data_structures::sync::Lrc;
@@ -1736,8 +1737,7 @@ rustc_queries! {
         arena_cache
         desc { |tcx| "resolving lifetimes for `{}`", tcx.def_path_str(def_id) }
     }
-    query named_variable_map(def_id: hir::OwnerId) ->
-        Option<&'tcx FxIndexMap<ItemLocalId, ResolvedArg>> {
+    query named_variable_map(def_id: hir::OwnerId) -> &'tcx SortedMap<ItemLocalId, ResolvedArg> {
         desc { |tcx| "looking up a named region inside `{}`", tcx.def_path_str(def_id) }
     }
     query is_late_bound_map(def_id: hir::OwnerId) -> Option<&'tcx FxIndexSet<ItemLocalId>> {
@@ -1753,7 +1753,7 @@ rustc_queries! {
         separate_provide_extern
     }
     query late_bound_vars_map(def_id: hir::OwnerId)
-        -> Option<&'tcx FxIndexMap<ItemLocalId, Vec<ty::BoundVariableKind>>> {
+        -> &'tcx SortedMap<ItemLocalId, Vec<ty::BoundVariableKind>> {
         desc { |tcx| "looking up late bound vars inside `{}`", tcx.def_path_str(def_id) }
     }
 
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -2988,7 +2988,7 @@ impl<'tcx> TyCtxt<'tcx> {
 
     pub fn named_bound_var(self, id: HirId) -> Option<resolve_bound_vars::ResolvedArg> {
         debug!(?id, "named_region");
-        self.named_variable_map(id.owner).and_then(|map| map.get(&id.local_id).cloned())
+        self.named_variable_map(id.owner).get(&id.local_id).cloned()
     }
 
     pub fn is_late_bound(self, id: HirId) -> bool {
@@ -2997,12 +2997,9 @@ impl<'tcx> TyCtxt<'tcx> {
 
     pub fn late_bound_vars(self, id: HirId) -> &'tcx List<ty::BoundVariableKind> {
         self.mk_bound_variable_kinds(
-            &self
-                .late_bound_vars_map(id.owner)
-                .and_then(|map| map.get(&id.local_id).cloned())
-                .unwrap_or_else(|| {
-                    bug!("No bound vars found for {}", self.hir().node_to_string(id))
-                }),
+            &self.late_bound_vars_map(id.owner).get(&id.local_id).cloned().unwrap_or_else(|| {
+                bug!("No bound vars found for {}", self.hir().node_to_string(id))
+            }),
         )
     }
 

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ use rustc_ast::expand::allocator::AllocatorKind;`
`16`	`16`	`use rustc_ast::expand::StrippedCfgItem;`
`17`	`17`	`use rustc_data_structures::fingerprint::Fingerprint;`
`18`	`18`	`use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};`
	`19`	`+use rustc_data_structures::sorted_map::SortedMap;`
`19`	`20`	`use rustc_data_structures::steal::Steal;`
`20`	`21`	`use rustc_data_structures::svh::Svh;`
`21`	`22`	`use rustc_data_structures::sync::Lrc;`
`@@ -1736,8 +1737,7 @@ rustc_queries! {`
`1736`	`1737`	`arena_cache`
`1737`	`1738`	desc { \|tcx\| "resolving lifetimes for `{}`", tcx.def_path_str(def_id) }
`1738`	`1739`	`}`
`1739`		`- query named_variable_map(def_id: hir::OwnerId) ->`
`1740`		`- Option<&'tcx FxIndexMap<ItemLocalId, ResolvedArg>> {`
	`1740`	`+ query named_variable_map(def_id: hir::OwnerId) -> &'tcx SortedMap<ItemLocalId, ResolvedArg> {`
`1741`	`1741`	desc { \|tcx\| "looking up a named region inside `{}`", tcx.def_path_str(def_id) }
`1742`	`1742`	`}`
`1743`	`1743`	`query is_late_bound_map(def_id: hir::OwnerId) -> Option<&'tcx FxIndexSet<ItemLocalId>> {`
`@@ -1753,7 +1753,7 @@ rustc_queries! {`
`1753`	`1753`	`separate_provide_extern`
`1754`	`1754`	`}`
`1755`	`1755`	`query late_bound_vars_map(def_id: hir::OwnerId)`
`1756`		`- -> Option<&'tcx FxIndexMap<ItemLocalId, Vec<ty::BoundVariableKind>>> {`
	`1756`	`+ -> &'tcx SortedMap<ItemLocalId, Vec<ty::BoundVariableKind>> {`
`1757`	`1757`	desc { \|tcx\| "looking up late bound vars inside `{}`", tcx.def_path_str(def_id) }
`1758`	`1758`	`}`
`1759`	`1759`