Auto merge of #90145 - cjgillot:sorted-map, r=michaelwoerister

Use SortedMap in HIR.

Closes https://github.com/rust-lang/rust/issues/89788
r? `@ghost`

3 files changed, 16 insertions, 29 deletions

diff --git a/compiler/rustc_ast_lowering/src/index.rs b/compiler/rustc_ast_lowering/src/index.rs
index dc2b1a730fb..8a9dad2cdd7 100644
--- a/compiler/rustc_ast_lowering/src/index.rs
+++ b/compiler/rustc_ast_lowering/src/index.rs
@@ -1,4 +1,5 @@
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
 use rustc_hir::definitions;
@@ -9,14 +10,13 @@ use rustc_session::Session;
 use rustc_span::source_map::SourceMap;
 use rustc_span::{Span, DUMMY_SP};
 
-use std::iter::repeat;
 use tracing::debug;
 
 /// A visitor that walks over the HIR and collects `Node`s into a HIR map.
 pub(super) struct NodeCollector<'a, 'hir> {
     /// Source map
     source_map: &'a SourceMap,
-    bodies: &'a IndexVec<ItemLocalId, Option<&'hir Body<'hir>>>,
+    bodies: &'a SortedMap<ItemLocalId, &'hir Body<'hir>>,
 
     /// Outputs
     nodes: IndexVec<ItemLocalId, Option<ParentedNode<'hir>>>,
@@ -30,21 +30,11 @@ pub(super) struct NodeCollector<'a, 'hir> {
     definitions: &'a definitions::Definitions,
 }
 
-fn insert_vec_map<K: Idx, V: Clone>(map: &mut IndexVec<K, Option<V>>, k: K, v: V) {
-    let i = k.index();
-    let len = map.len();
-    if i >= len {
-        map.extend(repeat(None).take(i - len + 1));
-    }
-    debug_assert!(map[k].is_none());
-    map[k] = Some(v);
-}
-
 pub(super) fn index_hir<'hir>(
     sess: &Session,
     definitions: &definitions::Definitions,
     item: hir::OwnerNode<'hir>,
-    bodies: &IndexVec<ItemLocalId, Option<&'hir Body<'hir>>>,
+    bodies: &SortedMap<ItemLocalId, &'hir Body<'hir>>,
 ) -> (IndexVec<ItemLocalId, Option<ParentedNode<'hir>>>, FxHashMap<LocalDefId, ItemLocalId>) {
     let mut nodes = IndexVec::new();
     // This node's parent should never be accessed: the owner's parent is computed by the
@@ -94,11 +84,7 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
             }
         }
 
-        insert_vec_map(
-            &mut self.nodes,
-            hir_id.local_id,
-            ParentedNode { parent: self.parent_node, node: node },
-        );
+        self.nodes.insert(hir_id.local_id, ParentedNode { parent: self.parent_node, node: node });
     }
 
     fn with_parent<F: FnOnce(&mut Self)>(&mut self, parent_node_id: HirId, f: F) {
@@ -144,7 +130,7 @@ impl<'a, 'hir> Visitor<'hir> for NodeCollector<'a, 'hir> {
 
     fn visit_nested_body(&mut self, id: BodyId) {
         debug_assert_eq!(id.hir_id.owner, self.owner);
-        let body = self.bodies[id.hir_id.local_id].unwrap();
+        let body = self.bodies[&id.hir_id.local_id];
         self.visit_body(body);
     }
 
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index e8747f2c5f8..63b20cd320b 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -974,8 +974,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let body = hir::Body { generator_kind: self.generator_kind, params, value };
         let id = body.id();
         debug_assert_eq!(id.hir_id.owner, self.current_hir_id_owner);
-        self.bodies.ensure_contains_elem(id.hir_id.local_id, || None);
-        self.bodies[id.hir_id.local_id] = Some(self.arena.alloc(body));
+        self.bodies.push((id.hir_id.local_id, self.arena.alloc(body)));
         id
     }
 
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 63ba5b45f4c..55cb4b035f6 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -45,6 +45,7 @@ use rustc_ast_pretty::pprust;
 use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::{struct_span_err, Applicability};
@@ -67,7 +68,6 @@ use rustc_span::symbol::{kw, sym, Ident, Symbol};
 use rustc_span::{Span, DUMMY_SP};
 
 use smallvec::SmallVec;
-use std::collections::BTreeMap;
 use tracing::{debug, trace};
 
 macro_rules! arena_vec {
@@ -104,9 +104,9 @@ struct LoweringContext<'a, 'hir: 'a> {
     /// The items being lowered are collected here.
     owners: IndexVec<LocalDefId, Option<hir::OwnerInfo<'hir>>>,
     /// Bodies inside the owner being lowered.
-    bodies: IndexVec<hir::ItemLocalId, Option<&'hir hir::Body<'hir>>>,
+    bodies: Vec<(hir::ItemLocalId, &'hir hir::Body<'hir>)>,
     /// Attributes inside the owner being lowered.
-    attrs: BTreeMap<hir::ItemLocalId, &'hir [Attribute]>,
+    attrs: SortedMap<hir::ItemLocalId, &'hir [Attribute]>,
 
     generator_kind: Option<hir::GeneratorKind>,
 
@@ -301,8 +301,8 @@ pub fn lower_crate<'a, 'hir>(
         nt_to_tokenstream,
         arena,
         owners,
-        bodies: IndexVec::new(),
-        attrs: BTreeMap::default(),
+        bodies: Vec::new(),
+        attrs: SortedMap::new(),
         catch_scope: None,
         loop_scope: None,
         is_in_loop_condition: false,
@@ -479,7 +479,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
     fn make_owner_info(&mut self, node: hir::OwnerNode<'hir>) -> hir::OwnerInfo<'hir> {
         let attrs = std::mem::take(&mut self.attrs);
-        let bodies = std::mem::take(&mut self.bodies);
+        let mut bodies = std::mem::take(&mut self.bodies);
         let local_node_ids = std::mem::take(&mut self.local_node_ids);
         let trait_map = local_node_ids
             .into_iter()
@@ -491,13 +491,15 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             .collect();
 
         #[cfg(debug_assertions)]
-        for (&id, attrs) in attrs.iter() {
+        for (id, attrs) in attrs.iter() {
             // Verify that we do not store empty slices in the map.
             if attrs.is_empty() {
                 panic!("Stored empty attributes for {:?}", id);
             }
         }
 
+        bodies.sort_by_key(|(k, _)| *k);
+        let bodies = SortedMap::from_presorted_elements(bodies);
         let (hash_including_bodies, hash_without_bodies) = self.hash_owner(node, &bodies);
         let (nodes, parenting) =
             index::index_hir(self.sess, self.resolver.definitions(), node, &bodies);
@@ -518,7 +520,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
     fn hash_owner(
         &mut self,
         node: hir::OwnerNode<'hir>,
-        bodies: &IndexVec<hir::ItemLocalId, Option<&'hir hir::Body<'hir>>>,
+        bodies: &SortedMap<hir::ItemLocalId, &'hir hir::Body<'hir>>,
     ) -> (Fingerprint, Fingerprint) {
         let mut hcx = self.resolver.create_stable_hashing_context();
         let mut stable_hasher = StableHasher::new();