Skip to content

Commit 396aeb8

Browse files
ZoxcZoxc
committed
Optimize the HIR map
1 parent aea57ae commit 396aeb8

File tree

3 files changed

+49
-46
lines changed

3 files changed

+49
-46
lines changed

‎src/librustc/hir/map/collector.rs‎

Lines changed: 36 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::arena::Arena;
22
use crate::hir::map::definitions::{self, DefPathHash};
3-
use crate::hir::map::{Entry, Map};
3+
use crate::hir::map::{Entry, HirOwnerData, Map};
44
use crate::hir::{HirItem, HirOwner, HirOwnerItems};
55
use crate::ich::StableHashingContext;
66
use crate::middle::cstore::CrateStore;
@@ -14,7 +14,7 @@ use rustc_hir::def_id::CRATE_DEF_INDEX;
1414
use rustc_hir::def_id::{CrateNum, DefIndex, LOCAL_CRATE};
1515
use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
1616
use rustc_hir::*;
17-
use rustc_index::vec::IndexVec;
17+
use rustc_index::vec::{Idx, IndexVec};
1818
use rustc_session::{CrateDisambiguator, Session};
1919
use rustc_span::source_map::SourceMap;
2020
use rustc_span::{Span, Symbol, DUMMY_SP};
@@ -31,8 +31,7 @@ pub(super) struct NodeCollector<'a, 'hir> {
3131
/// Source map
3232
source_map: &'a SourceMap,
3333

34-
owner_map: FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
35-
owner_items_map: FxHashMap<DefIndex, &'hir mut HirOwnerItems<'hir>>,
34+
map: IndexVec<DefIndex, HirOwnerData<'hir>>,
3635

3736
/// The parent of this node
3837
parent_node: hir::HirId,
@@ -49,6 +48,15 @@ pub(super) struct NodeCollector<'a, 'hir> {
4948
hir_body_nodes: Vec<(DefPathHash, Fingerprint)>,
5049
}
5150

51+
fn insert_vec_map<K: Idx, V: Clone>(map: &mut IndexVec<K, Option<V>>, k: K, v: V) {
52+
let i = k.index();
53+
let len = map.len();
54+
if i >= len {
55+
map.extend(repeat(None).take(i - len + 1));
56+
}
57+
map[k] = Some(v);
58+
}
59+
5260
fn hash(
5361
hcx: &mut StableHashingContext<'_>,
5462
input: impl for<'a> HashStable<StableHashingContext<'a>>,
@@ -126,14 +134,9 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
126134
hir_to_node_id,
127135
hcx,
128136
hir_body_nodes,
129-
owner_map: FxHashMap::with_capacity_and_hasher(
130-
definitions.def_index_count(),
131-
Default::default(),
132-
),
133-
owner_items_map: FxHashMap::with_capacity_and_hasher(
134-
definitions.def_index_count(),
135-
Default::default(),
136-
),
137+
map: (0..definitions.def_index_count())
138+
.map(|_| HirOwnerData { signature: None, with_bodies: None })
139+
.collect(),
137140
};
138141
collector.insert_entry(
139142
hir::CRATE_HIR_ID,
@@ -149,14 +152,10 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
149152
crate_disambiguator: CrateDisambiguator,
150153
cstore: &dyn CrateStore,
151154
commandline_args_hash: u64,
152-
) -> (
153-
FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
154-
FxHashMap<DefIndex, &'hir mut HirOwnerItems<'hir>>,
155-
Svh,
156-
) {
155+
) -> (IndexVec<DefIndex, HirOwnerData<'hir>>, Svh) {
157156
// Insert bodies into the map
158157
for (id, body) in self.krate.bodies.iter() {
159-
let bodies = &mut self.owner_items_map.get_mut(&id.hir_id.owner).unwrap().bodies;
158+
let bodies = &mut self.map[id.hir_id.owner].with_bodies.as_mut().unwrap().bodies;
160159
assert!(bodies.insert(id.hir_id.local_id, body).is_none());
161160
}
162161

@@ -196,39 +195,42 @@ impl<'a, 'hir> NodeCollector<'a, 'hir> {
196195
let crate_hash: Fingerprint = stable_hasher.finish();
197196

198197
let svh = Svh::new(crate_hash.to_smaller_hash());
199-
(self.owner_map, self.owner_items_map, svh)
198+
(self.map, svh)
200199
}
201200

202201
fn insert_entry(&mut self, id: HirId, entry: Entry<'hir>, hash: Fingerprint) {
203202
let i = id.local_id.as_u32() as usize;
204203

205-
let owner = HirOwner { parent: entry.parent, node: entry.node };
206-
207204
let arena = self.arena;
208205

209-
let items = self.owner_items_map.entry(id.owner).or_insert_with(|| {
210-
arena.alloc(HirOwnerItems {
206+
let data = &mut self.map[id.owner];
207+
208+
if data.with_bodies.is_none() {
209+
data.with_bodies = Some(arena.alloc(HirOwnerItems {
211210
hash,
212211
items: IndexVec::new(),
213212
bodies: FxHashMap::default(),
214-
})
215-
});
213+
}));
214+
}
215+
216+
let items = data.with_bodies.as_mut().unwrap();
216217

217218
if i == 0 {
218219
// Overwrite the dummy hash with the real HIR owner hash.
219220
items.hash = hash;
220221

221-
self.owner_map.insert(id.owner, self.arena.alloc(owner));
222-
// FIXME: feature(impl_trait_in_bindings) broken and trigger this assert
223-
//assert!(self.owner_map.insert(id.owner, self.arena.alloc(owner)).is_none());
222+
// FIXME: feature(impl_trait_in_bindings) broken and trigger this assert
223+
//assert!(data.signature.is_none());
224+
225+
data.signature =
226+
Some(self.arena.alloc(HirOwner { parent: entry.parent, node: entry.node }));
224227
} else {
225-
let len = items.items.len();
226-
if i >= len {
227-
items.items.extend(repeat(None).take(i - len + 1));
228-
}
229228
assert_eq!(entry.parent.owner, id.owner);
230-
items.items[id.local_id] =
231-
Some(HirItem { parent: entry.parent.local_id, node: entry.node });
229+
insert_vec_map(
230+
&mut items.items,
231+
id.local_id,
232+
HirItem { parent: entry.parent.local_id, node: entry.node },
233+
);
232234
}
233235
}
234236

‎src/librustc/hir/map/mod.rs‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ use rustc_hir::intravisit;
1515
use rustc_hir::itemlikevisit::ItemLikeVisitor;
1616
use rustc_hir::print::Nested;
1717
use rustc_hir::*;
18+
use rustc_index::vec::IndexVec;
1819
use rustc_span::hygiene::MacroKind;
1920
use rustc_span::source_map::Spanned;
2021
use rustc_span::symbol::kw;
@@ -128,12 +129,16 @@ fn is_body_owner<'hir>(node: Node<'hir>, hir_id: HirId) -> bool {
128129
}
129130
}
130131

132+
pub(super) struct HirOwnerData<'hir> {
133+
pub(super) signature: Option<&'hir HirOwner<'hir>>,
134+
pub(super) with_bodies: Option<&'hir mut HirOwnerItems<'hir>>,
135+
}
136+
131137
pub struct IndexedHir<'hir> {
132138
/// The SVH of the local crate.
133139
pub crate_hash: Svh,
134140

135-
pub(super) owner_map: FxHashMap<DefIndex, &'hir HirOwner<'hir>>,
136-
pub(super) owner_items_map: FxHashMap<DefIndex, &'hir HirOwnerItems<'hir>>,
141+
pub(super) map: IndexVec<DefIndex, HirOwnerData<'hir>>,
137142

138143
/// The reverse mapping of `node_to_hir_id`.
139144
pub(super) hir_to_node_id: FxHashMap<HirId, NodeId>,
@@ -1036,7 +1041,7 @@ pub(super) fn index_hir<'tcx>(tcx: TyCtxt<'tcx>, cnum: CrateNum) -> &'tcx Indexe
10361041
.map(|(node_id, &hir_id)| (hir_id, node_id))
10371042
.collect();
10381043

1039-
let (owner_map, owner_items_map, crate_hash) = {
1044+
let (map, crate_hash) = {
10401045
let hcx = tcx.create_stable_hashing_context();
10411046

10421047
let mut collector = NodeCollector::root(
@@ -1054,12 +1059,7 @@ pub(super) fn index_hir<'tcx>(tcx: TyCtxt<'tcx>, cnum: CrateNum) -> &'tcx Indexe
10541059
collector.finalize_and_compute_crate_hash(crate_disambiguator, &*tcx.cstore, cmdline_args)
10551060
};
10561061

1057-
let map = tcx.arena.alloc(IndexedHir {
1058-
crate_hash,
1059-
owner_map,
1060-
owner_items_map: owner_items_map.into_iter().map(|(k, v)| (k, &*v)).collect(),
1061-
hir_to_node_id,
1062-
});
1062+
let map = tcx.arena.alloc(IndexedHir { crate_hash, map, hir_to_node_id });
10631063

10641064
map
10651065
}

‎src/librustc/hir/mod.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,9 @@ pub fn provide(providers: &mut Providers<'_>) {
7878
let module = hir.as_local_hir_id(id).unwrap();
7979
&tcx.untracked_crate.modules[&module]
8080
};
81-
providers.hir_owner = |tcx, id| *tcx.index_hir(id.krate).owner_map.get(&id.index).unwrap();
82-
providers.hir_owner_items =
83-
|tcx, id| *tcx.index_hir(id.krate).owner_items_map.get(&id.index).unwrap();
81+
providers.hir_owner = |tcx, id| tcx.index_hir(id.krate).map[id.index].signature.unwrap();
82+
providers.hir_owner_items = |tcx, id| {
83+
tcx.index_hir(id.krate).map[id.index].with_bodies.as_ref().map(|items| &**items).unwrap()
84+
};
8485
map::provide(providers);
8586
}

0 commit comments

Comments
 (0)