Use an IndexVec for bodies.

Author: cjgillot

compiler/rustc_ast_lowering/src/item.rs

@@ -974,7 +974,8 @@ 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.insert(id.hir_id.local_id, body);
+        self.bodies.ensure_contains_elem(id.hir_id.local_id, || None);
+        self.bodies[id.hir_id.local_id] = Some(self.arena.alloc(body));
         id
     }
 

compiler/rustc_ast_lowering/src/lib.rs

@@ -97,7 +97,7 @@ struct LoweringContext<'a, 'hir: 'a> {
 
     /// The items being lowered are collected here.
     owners: IndexVec<LocalDefId, Option<hir::OwnerInfo<'hir>>>,
-    bodies: BTreeMap<hir::ItemLocalId, hir::Body<'hir>>,
+    bodies: IndexVec<hir::ItemLocalId, Option<&'hir hir::Body<'hir>>>,
     attrs: BTreeMap<hir::ItemLocalId, &'hir [Attribute]>,
 
     generator_kind: Option<hir::GeneratorKind>,
@@ -322,7 +322,7 @@ pub fn lower_crate<'a, 'hir>(
         nt_to_tokenstream,
         arena,
         owners,
-        bodies: BTreeMap::new(),
+        bodies: IndexVec::new(),
         attrs: BTreeMap::default(),
         catch_scope: None,
         loop_scope: None,

compiler/rustc_hir/src/arena.rs

@@ -19,6 +19,7 @@ macro_rules! arena_types {
             [] attribute: rustc_ast::Attribute,
             [] block: rustc_hir::Block<$tcx>,
             [] bare_fn_ty: rustc_hir::BareFnTy<$tcx>,
+            [] body: rustc_hir::Body<$tcx>,
             [] generic_arg: rustc_hir::GenericArg<$tcx>,
             [] generic_args: rustc_hir::GenericArgs<$tcx>,
             [] generic_bound: rustc_hir::GenericBound<$tcx>,

compiler/rustc_hir/src/hir.rs

@@ -666,7 +666,7 @@ pub struct WhereEqPredicate<'hir> {
 pub struct OwnerInfo<'hir> {
     pub node: OwnerNode<'hir>,
     pub attrs: BTreeMap<ItemLocalId, &'hir [Attribute]>,
-    pub bodies: BTreeMap<ItemLocalId, Body<'hir>>,
+    pub bodies: IndexVec<ItemLocalId, Option<&'hir Body<'hir>>>,
     /// Map indicating what traits are in scope for places where this
     /// is relevant; generated by resolve.
     pub trait_map: FxHashMap<ItemLocalId, Box<[TraitCandidate]>>,
@@ -705,9 +705,9 @@ impl Crate<'hir> {
         self.owners[id.def_id].as_ref().unwrap().node.expect_foreign_item()
     }
 
-    pub fn body(&self, id: BodyId) -> &Body<'hir> {
+    pub fn body(&self, id: BodyId) -> &'hir Body<'hir> {
         let HirId { owner, local_id } = id.hir_id;
-        &self.owners[owner].as_ref().unwrap().bodies[&local_id]
+        self.owners[owner].as_ref().unwrap().bodies[local_id].unwrap()
     }
 
     pub fn attrs(&self, id: HirId) -> &'hir [Attribute] {

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

@@ -96,11 +96,7 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
         let mut nodes = IndexVec::new();
         nodes.push(Some(ParentedNode { parent: ItemLocalId::new(0), node: node.into() }));
 
-        let mut bodies = FxHashMap::default();
-        for (id, body) in self.krate.owners[owner].as_ref().unwrap().bodies.iter() {
-            let _old = bodies.insert(*id, body);
-            debug_assert!(_old.is_none());
-        }
+        let bodies = &self.krate.owners[owner].as_ref().unwrap().bodies;
 
         debug_assert!(self.map[owner].is_none());
         self.map[owner] = Some(self.arena.alloc(OwnerNodes { hash, nodes, bodies }));

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

@@ -381,7 +381,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.get(&id.hir_id.local_id).unwrap()
+        self.tcx.hir_owner_nodes(id.hir_id.owner).unwrap().bodies[id.hir_id.local_id].unwrap()
     }
 
     pub fn fn_decl_by_hir_id(&self, hir_id: HirId) -> Option<&'hir FnDecl<'hir>> {
@@ -500,10 +500,13 @@ impl<'hir> Map<'hir> {
             .iter_enumerated()
             .flat_map(move |(owner, owner_info)| {
                 let bodies = &owner_info.as_ref()?.bodies;
-                Some(bodies.keys().map(move |&local_id| {
+                Some(bodies.iter_enumerated().filter_map(move |(local_id, body)| {
+                    if body.is_none() {
+                        return None;
+                    }
                     let hir_id = HirId { owner, local_id };
                     let body_id = BodyId { hir_id };
-                    self.body_owner_def_id(body_id)
+                    Some(self.body_owner_def_id(body_id))
                 }))
             })
             .flatten()
@@ -517,10 +520,13 @@ 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.bodies).for_each(|(&local_id, _)| {
-                    let hir_id = HirId { owner, local_id };
-                    let body_id = BodyId { hir_id };
-                    f(self.body_owner_def_id(body_id))
+                par_iter(&owner_info.bodies.raw).enumerate().for_each(|(local_id, body)| {
+                    if body.is_some() {
+                        let local_id = ItemLocalId::new(local_id);
+                        let hir_id = HirId { owner, local_id };
+                        let body_id = BodyId { hir_id };
+                        f(self.body_owner_def_id(body_id))
+                    }
                 })
             }
         });

compiler/rustc_middle/src/hir/mod.rs

@@ -65,7 +65,7 @@ pub struct OwnerNodes<'tcx> {
     // The zeroth node's parent is trash, but is never accessed.
     nodes: IndexVec<ItemLocalId, Option<ParentedNode<'tcx>>>,
     /// Content of local bodies.
-    bodies: FxHashMap<ItemLocalId, &'tcx Body<'tcx>>,
+    bodies: &'tcx IndexVec<ItemLocalId, Option<&'tcx Body<'tcx>>>,
 }
 
 impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for OwnerNodes<'tcx> {