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

Use SortedMap in HIR.

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

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(
+        self.nodes.insert(hir_id.local_id, ParentedNode { parent: self.parent_node, node: node });
-            &mut self.nodes,
-            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);
     }
 

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.push((id.hir_id.local_id, self.arena.alloc(body)));
-        self.bodies[id.hir_id.local_id] = Some(self.arena.alloc(body));
         id
     }
 

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(),
+        bodies: Vec::new(),
-        attrs: BTreeMap::default(),
+        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();

compiler/rustc_data_structures/src/sorted_map.rs

@@ -1,3 +1,4 @@
+use crate::stable_hasher::{HashStable, StableHasher};
 use std::borrow::Borrow;
 use std::cmp::Ordering;
 use std::iter::FromIterator;
@@ -16,17 +17,26 @@ pub use index_map::SortedIndexMultiMap;
 /// stores data in a more compact way. It also supports accessing contiguous
 /// ranges of elements as a slice, and slices of already sorted elements can be
 /// inserted efficiently.
-#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default, Debug, Encodable, Decodable)]
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Encodable, Decodable)]
-pub struct SortedMap<K: Ord, V> {
+pub struct SortedMap<K, V> {
     data: Vec<(K, V)>,
 }
 
-impl<K: Ord, V> SortedMap<K, V> {
+impl<K, V> Default for SortedMap<K, V> {
     #[inline]
-    pub fn new() -> SortedMap<K, V> {
+    fn default() -> SortedMap<K, V> {
-        SortedMap { data: vec![] }
+        SortedMap { data: Vec::new() }
     }
+}
 
+impl<K, V> SortedMap<K, V> {
+    #[inline]
+    pub const fn new() -> SortedMap<K, V> {
+        SortedMap { data: Vec::new() }
+    }
+}
+
+impl<K: Ord, V> SortedMap<K, V> {
     /// Construct a `SortedMap` from a presorted set of elements. This is faster
     /// than creating an empty map and then inserting the elements individually.
     ///
@@ -281,5 +291,12 @@ impl<K: Ord, V> FromIterator<(K, V)> for SortedMap<K, V> {
     }
 }
 
+impl<K: HashStable<CTX>, V: HashStable<CTX>, CTX> HashStable<CTX> for SortedMap<K, V> {
+    #[inline]
+    fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
+        self.data.hash_stable(ctx, hasher);
+    }
+}
+
 #[cfg(test)]
 mod tests;

compiler/rustc_hir/src/hir.rs

@@ -12,6 +12,7 @@ pub use rustc_ast::{CaptureBy, Movability, Mutability};
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_index::vec::IndexVec;
 use rustc_macros::HashStable_Generic;
 use rustc_span::source_map::Spanned;
@@ -22,7 +23,6 @@ use rustc_target::asm::InlineAsmRegOrRegClass;
 use rustc_target::spec::abi::Abi;
 
 use smallvec::SmallVec;
-use std::collections::BTreeMap;
 use std::fmt;
 
 #[derive(Copy, Clone, Encodable, HashStable_Generic)]
@@ -676,13 +676,13 @@ pub struct ParentedNode<'tcx> {
 /// Attributes owned by a HIR owner.
 #[derive(Debug)]
 pub struct AttributeMap<'tcx> {
-    pub map: BTreeMap<ItemLocalId, &'tcx [Attribute]>,
+    pub map: SortedMap<ItemLocalId, &'tcx [Attribute]>,
     pub hash: Fingerprint,
 }
 
 impl<'tcx> AttributeMap<'tcx> {
     pub const EMPTY: &'static AttributeMap<'static> =
-        &AttributeMap { map: BTreeMap::new(), hash: Fingerprint::ZERO };
+        &AttributeMap { map: SortedMap::new(), hash: Fingerprint::ZERO };
 
     #[inline]
     pub fn get(&self, id: ItemLocalId) -> &'tcx [Attribute] {
@@ -705,7 +705,7 @@ pub struct OwnerNodes<'tcx> {
     // used.
     pub nodes: IndexVec<ItemLocalId, Option<ParentedNode<'tcx>>>,
     /// Content of local bodies.
-    pub bodies: IndexVec<ItemLocalId, Option<&'tcx Body<'tcx>>>,
+    pub bodies: SortedMap<ItemLocalId, &'tcx Body<'tcx>>,
 }
 
 /// Full information resulting from lowering an AST node.

compiler/rustc_middle/src/hir/map/mod.rs

@@ -376,7 +376,7 @@ impl<'hir> Map<'hir> {
     }
 
     pub fn body(&self, id: BodyId) -> &'hir Body<'hir> {
-        self.tcx.hir_owner_nodes(id.hir_id.owner).unwrap().bodies[id.hir_id.local_id].unwrap()
+        self.tcx.hir_owner_nodes(id.hir_id.owner).unwrap().bodies[&id.hir_id.local_id]
     }
 
     pub fn fn_decl_by_hir_id(&self, hir_id: HirId) -> Option<&'hir FnDecl<'hir>> {
@@ -495,13 +495,10 @@ impl<'hir> Map<'hir> {
             .iter_enumerated()
             .flat_map(move |(owner, owner_info)| {
                 let bodies = &owner_info.as_ref()?.nodes.bodies;
-                Some(bodies.iter_enumerated().filter_map(move |(local_id, body)| {
+                Some(bodies.iter().map(move |&(local_id, _)| {
-                    if body.is_none() {
-                        return None;
-                    }
                     let hir_id = HirId { owner, local_id };
                     let body_id = BodyId { hir_id };
-                    Some(self.body_owner_def_id(body_id))
+                    self.body_owner_def_id(body_id)
                 }))
             })
             .flatten()
@@ -515,13 +512,10 @@ impl<'hir> Map<'hir> {
         par_iter(&self.krate().owners.raw).enumerate().for_each(|(owner, owner_info)| {
             let owner = LocalDefId::new(owner);
             if let Some(owner_info) = owner_info {
-                par_iter(&owner_info.nodes.bodies.raw).enumerate().for_each(|(local_id, body)| {
+                par_iter(owner_info.nodes.bodies.range(..)).for_each(|(local_id, _)| {
-                    if body.is_some() {
+                    let hir_id = HirId { owner, local_id: *local_id };
-                        let local_id = ItemLocalId::new(local_id);
+                    let body_id = BodyId { hir_id };
-                        let hir_id = HirId { owner, local_id };
+                    f(self.body_owner_def_id(body_id))
-                        let body_id = BodyId { hir_id };
-                        f(self.body_owner_def_id(body_id))
-                    }
                 })
             }
         });
@@ -578,8 +572,8 @@ impl<'hir> Map<'hir> {
         let krate = self.krate();
         for (owner, info) in krate.owners.iter_enumerated() {
             if let Some(info) = info {
-                for (&local_id, attrs) in info.attrs.map.iter() {
+                for (local_id, attrs) in info.attrs.map.iter() {
-                    let id = HirId { owner, local_id };
+                    let id = HirId { owner, local_id: *local_id };
                     for a in *attrs {
                         visitor.visit_attribute(id, a)
                     }

compiler/rustc_query_system/src/ich/hcx.rs

@@ -1,12 +1,12 @@
 use crate::ich;
 use rustc_ast as ast;
 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_hir as hir;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::definitions::{DefPathHash, Definitions};
-use rustc_index::vec::IndexVec;
 use rustc_session::cstore::CrateStore;
 use rustc_session::Session;
 use rustc_span::source_map::SourceMap;
@@ -51,7 +51,7 @@ pub(super) enum BodyResolver<'tcx> {
     Traverse {
         hash_bodies: bool,
         owner: LocalDefId,
-        bodies: &'tcx IndexVec<hir::ItemLocalId, Option<&'tcx hir::Body<'tcx>>>,
+        bodies: &'tcx SortedMap<hir::ItemLocalId, &'tcx hir::Body<'tcx>>,
     },
 }
 
@@ -122,7 +122,7 @@ impl<'a> StableHashingContext<'a> {
         &mut self,
         hash_bodies: bool,
         owner: LocalDefId,
-        bodies: &'a IndexVec<hir::ItemLocalId, Option<&'a hir::Body<'a>>>,
+        bodies: &'a SortedMap<hir::ItemLocalId, &'a hir::Body<'a>>,
         f: impl FnOnce(&mut Self),
     ) {
         let prev = self.body_resolver;

compiler/rustc_query_system/src/ich/impls_hir.rs

@@ -33,7 +33,7 @@ impl<'ctx> rustc_hir::HashStableContext for StableHashingContext<'ctx> {
             BodyResolver::Traverse { hash_bodies: false, .. } => {}
             BodyResolver::Traverse { hash_bodies: true, owner, bodies } => {
                 assert_eq!(id.hir_id.owner, owner);
-                bodies[id.hir_id.local_id].unwrap().hash_stable(hcx, hasher);
+                bodies[&id.hir_id.local_id].hash_stable(hcx, hasher);
             }
         }
     }