Skip to content

Commit 42c9514

Browse files
cjgillotcjgillot
committed
Simplify construction of replacement map.
1 parent e465d64 commit 42c9514

11 files changed

+196
-160
lines changed

‎compiler/rustc_mir_dataflow/src/value_analysis.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -790,7 +790,7 @@ impl<V, T> TryFrom<ProjectionElem<V, T>> for TrackElem {
790790
}
791791

792792
/// Invokes `f` on all direct fields of `ty`.
793-
fn iter_fields<'tcx>(
793+
pub fn iter_fields<'tcx>(
794794
ty: Ty<'tcx>,
795795
tcx: TyCtxt<'tcx>,
796796
mut f: impl FnMut(Option<VariantIdx>, Field, Ty<'tcx>),

‎compiler/rustc_mir_transform/src/sroa.rs‎

Lines changed: 69 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
use crate::MirPass;
2-
use rustc_data_structures::fx::{FxIndexMap, IndexEntry};
2+
use rustc_data_structures::fx::FxIndexMap;
33
use rustc_index::bit_set::BitSet;
44
use rustc_index::vec::IndexVec;
55
use rustc_middle::mir::patch::MirPatch;
66
use rustc_middle::mir::visit::*;
77
use rustc_middle::mir::*;
88
use rustc_middle::ty::TyCtxt;
9+
use rustc_mir_dataflow::value_analysis::iter_fields;
910

1011
pub struct ScalarReplacementOfAggregates;
1112

@@ -125,6 +126,36 @@ fn escaping_locals(body: &Body<'_>) -> BitSet<Local> {
125126
#[derive(Default, Debug)]
126127
struct ReplacementMap<'tcx> {
127128
fields: FxIndexMap<PlaceRef<'tcx>, Local>,
129+
/// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage
130+
/// and deinit statement and debuginfo.
131+
fragments: IndexVec<Local, Option<Vec<(&'tcx [PlaceElem<'tcx>], Local)>>>,
132+
}
133+
134+
impl<'tcx> ReplacementMap<'tcx> {
135+
fn gather_debug_info_fragments(
136+
&self,
137+
place: PlaceRef<'tcx>,
138+
) -> Option<Vec<VarDebugInfoFragment<'tcx>>> {
139+
let mut fragments = Vec::new();
140+
let Some(parts) = &self.fragments[place.local] else { return None };
141+
for (proj, replacement_local) in parts {
142+
if proj.starts_with(place.projection) {
143+
fragments.push(VarDebugInfoFragment {
144+
projection: proj[place.projection.len()..].to_vec(),
145+
contents: Place::from(*replacement_local),
146+
});
147+
}
148+
}
149+
Some(fragments)
150+
}
151+
152+
fn place_fragments(
153+
&self,
154+
place: Place<'tcx>,
155+
) -> Option<&Vec<(&'tcx [PlaceElem<'tcx>], Local)>> {
156+
let local = place.as_local()?;
157+
self.fragments[local].as_ref()
158+
}
128159
}
129160

130161
/// Compute the replacement of flattened places into locals.
@@ -136,53 +167,30 @@ fn compute_flattening<'tcx>(
136167
body: &mut Body<'tcx>,
137168
escaping: BitSet<Local>,
138169
) -> ReplacementMap<'tcx> {
139-
let mut visitor = PreFlattenVisitor {
140-
tcx,
141-
escaping,
142-
local_decls: &mut body.local_decls,
143-
map: Default::default(),
144-
};
145-
for (block, bbdata) in body.basic_blocks.iter_enumerated() {
146-
visitor.visit_basic_block_data(block, bbdata);
147-
}
148-
return visitor.map;
149-
150-
struct PreFlattenVisitor<'tcx, 'll> {
151-
tcx: TyCtxt<'tcx>,
152-
local_decls: &'ll mut LocalDecls<'tcx>,
153-
escaping: BitSet<Local>,
154-
map: ReplacementMap<'tcx>,
155-
}
170+
let mut fields = FxIndexMap::default();
171+
let mut fragments = IndexVec::from_elem(None::<Vec<_>>, &body.local_decls);
156172

157-
impl<'tcx, 'll> PreFlattenVisitor<'tcx, 'll> {
158-
fn create_place(&mut self, place: PlaceRef<'tcx>) {
159-
if self.escaping.contains(place.local) {
160-
return;
161-
}
162-
163-
match self.map.fields.entry(place) {
164-
IndexEntry::Occupied(_) => {}
165-
IndexEntry::Vacant(v) => {
166-
let ty = place.ty(&*self.local_decls, self.tcx).ty;
167-
let local = self.local_decls.push(LocalDecl {
168-
ty,
169-
user_ty: None,
170-
..self.local_decls[place.local].clone()
171-
});
172-
v.insert(local);
173-
}
174-
}
175-
}
176-
}
177-
178-
impl<'tcx, 'll> Visitor<'tcx> for PreFlattenVisitor<'tcx, 'll> {
179-
fn visit_place(&mut self, place: &Place<'tcx>, _: PlaceContext, _: Location) {
180-
if let &[PlaceElem::Field(..), ..] = &place.projection[..] {
181-
let pr = PlaceRef { local: place.local, projection: &place.projection[..1] };
182-
self.create_place(pr)
183-
}
173+
for local in body.local_decls.indices() {
174+
if escaping.contains(local) {
175+
continue;
184176
}
177+
let decl = body.local_decls[local].clone();
178+
let ty = decl.ty;
179+
iter_fields(ty, tcx, |variant, field, field_ty| {
180+
if variant.is_some() {
181+
// Downcasts are currently not supported.
182+
return;
183+
};
184+
let new_local =
185+
body.local_decls.push(LocalDecl { ty: field_ty, user_ty: None, ..decl.clone() });
186+
let place = Place::from(local)
187+
.project_deeper(&[PlaceElem::Field(field, field_ty)], tcx)
188+
.as_ref();
189+
fields.insert(place, new_local);
190+
fragments[local].get_or_insert_default().push((place.projection, new_local));
191+
});
185192
}
193+
ReplacementMap { fields, fragments }
186194
}
187195

188196
/// Perform the replacement computed by `compute_flattening`.
@@ -200,18 +208,11 @@ fn replace_flattened_locals<'tcx>(
200208
return;
201209
}
202210

203-
let mut fragments = IndexVec::<_, Option<Vec<_>>>::from_elem(None, &body.local_decls);
204-
for (k, v) in &replacements.fields {
205-
fragments[k.local].get_or_insert_default().push((k.projection, *v));
206-
}
207-
debug!(?fragments);
208-
209211
let mut visitor = ReplacementVisitor {
210212
tcx,
211213
local_decls: &body.local_decls,
212214
replacements,
213215
all_dead_locals,
214-
fragments: &fragments,
215216
patch: MirPatch::new(body),
216217
};
217218
for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
@@ -237,30 +238,10 @@ struct ReplacementVisitor<'tcx, 'll> {
237238
replacements: ReplacementMap<'tcx>,
238239
/// This is used to check that we are not leaving references to replaced locals behind.
239240
all_dead_locals: BitSet<Local>,
240-
/// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage
241-
/// and deinit statement and debuginfo.
242-
fragments: &'ll IndexVec<Local, Option<Vec<(&'tcx [PlaceElem<'tcx>], Local)>>>,
243241
patch: MirPatch<'tcx>,
244242
}
245243

246244
impl<'tcx, 'll> ReplacementVisitor<'tcx, 'll> {
247-
fn gather_debug_info_fragments(
248-
&self,
249-
place: PlaceRef<'tcx>,
250-
) -> Option<Vec<VarDebugInfoFragment<'tcx>>> {
251-
let mut fragments = Vec::new();
252-
let Some(parts) = &self.fragments[place.local] else { return None };
253-
for (proj, replacement_local) in parts {
254-
if proj.starts_with(place.projection) {
255-
fragments.push(VarDebugInfoFragment {
256-
projection: proj[place.projection.len()..].to_vec(),
257-
contents: Place::from(*replacement_local),
258-
});
259-
}
260-
}
261-
Some(fragments)
262-
}
263-
264245
fn replace_place(&self, place: PlaceRef<'tcx>) -> Option<Place<'tcx>> {
265246
if let &[PlaceElem::Field(..), ref rest @ ..] = place.projection {
266247
let pr = PlaceRef { local: place.local, projection: &place.projection[..1] };
@@ -270,25 +251,18 @@ impl<'tcx, 'll> ReplacementVisitor<'tcx, 'll> {
270251
None
271252
}
272253
}
273-
274-
fn place_fragments(
275-
&self,
276-
place: Place<'tcx>,
277-
) -> Option<&'ll Vec<(&'tcx [PlaceElem<'tcx>], Local)>> {
278-
let local = place.as_local()?;
279-
self.fragments[local].as_ref()
280-
}
281254
}
282255

283256
impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
284257
fn tcx(&self) -> TyCtxt<'tcx> {
285258
self.tcx
286259
}
287260

261+
#[instrument(level = "trace", skip(self))]
288262
fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
289263
match statement.kind {
290264
StatementKind::StorageLive(l) => {
291-
if let Some(final_locals) = &self.fragments[l] {
265+
if let Some(final_locals) = &self.replacements.fragments[l] {
292266
for &(_, fl) in final_locals {
293267
self.patch.add_statement(location, StatementKind::StorageLive(fl));
294268
}
@@ -297,7 +271,7 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
297271
return;
298272
}
299273
StatementKind::StorageDead(l) => {
300-
if let Some(final_locals) = &self.fragments[l] {
274+
if let Some(final_locals) = &self.replacements.fragments[l] {
301275
for &(_, fl) in final_locals {
302276
self.patch.add_statement(location, StatementKind::StorageDead(fl));
303277
}
@@ -306,7 +280,7 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
306280
return;
307281
}
308282
StatementKind::Deinit(box place) => {
309-
if let Some(final_locals) = self.place_fragments(place) {
283+
if let Some(final_locals) = self.replacements.place_fragments(place) {
310284
for &(_, fl) in final_locals {
311285
self.patch
312286
.add_statement(location, StatementKind::Deinit(Box::new(fl.into())));
@@ -317,7 +291,7 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
317291
}
318292

319293
StatementKind::Assign(box (place, Rvalue::Aggregate(_, ref operands))) => {
320-
if let Some(final_locals) = self.place_fragments(place) {
294+
if let Some(final_locals) = self.replacements.place_fragments(place) {
321295
for &(projection, fl) in final_locals {
322296
let &[PlaceElem::Field(index, _)] = projection else { bug!() };
323297
let index = index.as_usize();
@@ -333,7 +307,7 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
333307
}
334308

335309
StatementKind::Assign(box (place, Rvalue::Use(Operand::Constant(_)))) => {
336-
if let Some(final_locals) = self.place_fragments(place) {
310+
if let Some(final_locals) = self.replacements.place_fragments(place) {
337311
for &(projection, fl) in final_locals {
338312
let rvalue =
339313
Rvalue::Use(Operand::Move(place.project_deeper(projection, self.tcx)));
@@ -353,9 +327,12 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
353327
Operand::Move(rplace) => (rplace, false),
354328
Operand::Constant(_) => bug!(),
355329
};
356-
if let Some(final_locals) = self.place_fragments(lhs) {
330+
if let Some(final_locals) = self.replacements.place_fragments(lhs) {
357331
for &(projection, fl) in final_locals {
358332
let rplace = rplace.project_deeper(projection, self.tcx);
333+
debug!(?rplace);
334+
let rplace = self.replace_place(rplace.as_ref()).unwrap_or(rplace);
335+
debug!(?rplace);
359336
let rvalue = if copy {
360337
Rvalue::Use(Operand::Copy(rplace))
361338
} else {
@@ -389,7 +366,9 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
389366
VarDebugInfoContents::Place(ref mut place) => {
390367
if let Some(repl) = self.replace_place(place.as_ref()) {
391368
*place = repl;
392-
} else if let Some(fragments) = self.gather_debug_info_fragments(place.as_ref()) {
369+
} else if let Some(fragments) =
370+
self.replacements.gather_debug_info_fragments(place.as_ref())
371+
{
393372
let ty = place.ty(self.local_decls, self.tcx).ty;
394373
var_debug_info.value = VarDebugInfoContents::Composite { ty, fragments };
395374
}
@@ -401,8 +380,9 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
401380
if let Some(repl) = self.replace_place(fragment.contents.as_ref()) {
402381
fragment.contents = repl;
403382
true
404-
} else if let Some(frg) =
405-
self.gather_debug_info_fragments(fragment.contents.as_ref())
383+
} else if let Some(frg) = self
384+
.replacements
385+
.gather_debug_info_fragments(fragment.contents.as_ref())
406386
{
407387
new_fragments.extend(frg.into_iter().map(|mut f| {
408388
f.projection.splice(0..0, fragment.projection.iter().copied());

‎tests/mir-opt/const_debuginfo.main.ConstDebugInfo.diff‎

Lines changed: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88
let mut _6: u8; // in scope 0 at $DIR/const_debuginfo.rs:+4:15: +4:16
99
let mut _7: u8; // in scope 0 at $DIR/const_debuginfo.rs:+4:19: +4:20
1010
let mut _8: u8; // in scope 0 at $DIR/const_debuginfo.rs:+4:23: +4:24
11-
let mut _13: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:13: +13:16
12-
let mut _14: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:19: +13:22
11+
let mut _12: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:13: +13:16
12+
let mut _13: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:19: +13:22
1313
scope 1 {
1414
- debug x => _1; // in scope 1 at $DIR/const_debuginfo.rs:+1:9: +1:10
1515
+ debug x => const 1_u8; // in scope 1 at $DIR/const_debuginfo.rs:+1:9: +1:10
@@ -29,19 +29,21 @@
2929
scope 5 {
3030
- debug s => _9; // in scope 5 at $DIR/const_debuginfo.rs:+6:9: +6:10
3131
+ debug s => const "hello, world!"; // in scope 5 at $DIR/const_debuginfo.rs:+6:9: +6:10
32-
let _10: (bool, bool, u32); // in scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
32+
let _14: bool; // in scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
33+
let _15: bool; // in scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
34+
let _16: u32; // in scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
3335
scope 6 {
34-
debug f => _10; // in scope 6 at $DIR/const_debuginfo.rs:+8:9: +8:10
35-
let _11: std::option::Option<u16>; // in scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
36+
debug f => (bool, bool, u32){ .0 => _14, .1 => _15, .2 => _16, }; // in scope 6 at $DIR/const_debuginfo.rs:+8:9: +8:10
37+
let _10: std::option::Option<u16>; // in scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
3638
scope 7 {
37-
debug o => _11; // in scope 7 at $DIR/const_debuginfo.rs:+10:9: +10:10
38-
let _15: u32; // in scope 7 at $DIR/const_debuginfo.rs:+12:9: +12:10
39-
let _16: u32; // in scope 7 at $DIR/const_debuginfo.rs:+12:9: +12:10
39+
debug o => _10; // in scope 7 at $DIR/const_debuginfo.rs:+10:9: +10:10
40+
let _17: u32; // in scope 7 at $DIR/const_debuginfo.rs:+12:9: +12:10
41+
let _18: u32; // in scope 7 at $DIR/const_debuginfo.rs:+12:9: +12:10
4042
scope 8 {
41-
debug p => Point{ .0 => _15, .1 => _16, }; // in scope 8 at $DIR/const_debuginfo.rs:+12:9: +12:10
42-
let _12: u32; // in scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
43+
debug p => Point{ .0 => _17, .1 => _18, }; // in scope 8 at $DIR/const_debuginfo.rs:+12:9: +12:10
44+
let _11: u32; // in scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
4345
scope 9 {
44-
- debug a => _12; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
46+
- debug a => _11; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
4547
+ debug a => const 64_u32; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
4648
}
4749
}
@@ -67,17 +69,23 @@
6769
// mir::Constant
6870
// + span: $DIR/const_debuginfo.rs:14:13: 14:28
6971
// + literal: Const { ty: &str, val: Value(Slice(..)) }
70-
StorageLive(_10); // scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
71-
_10 = (const true, const false, const 123_u32); // scope 5 at $DIR/const_debuginfo.rs:+8:13: +8:34
72-
StorageLive(_11); // scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
73-
_11 = Option::<u16>::Some(const 99_u16); // scope 6 at $DIR/const_debuginfo.rs:+10:13: +10:24
74-
_15 = const 32_u32; // scope 7 at $DIR/const_debuginfo.rs:+12:13: +12:35
75-
_16 = const 32_u32; // scope 7 at $DIR/const_debuginfo.rs:+12:13: +12:35
76-
StorageLive(_12); // scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
77-
_12 = const 64_u32; // scope 8 at $DIR/const_debuginfo.rs:+13:13: +13:22
78-
StorageDead(_12); // scope 8 at $DIR/const_debuginfo.rs:+14:1: +14:2
79-
StorageDead(_11); // scope 6 at $DIR/const_debuginfo.rs:+14:1: +14:2
80-
StorageDead(_10); // scope 5 at $DIR/const_debuginfo.rs:+14:1: +14:2
72+
StorageLive(_14); // scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
73+
StorageLive(_15); // scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
74+
StorageLive(_16); // scope 5 at $DIR/const_debuginfo.rs:+8:9: +8:10
75+
_14 = const true; // scope 5 at $DIR/const_debuginfo.rs:+8:13: +8:34
76+
_15 = const false; // scope 5 at $DIR/const_debuginfo.rs:+8:13: +8:34
77+
_16 = const 123_u32; // scope 5 at $DIR/const_debuginfo.rs:+8:13: +8:34
78+
StorageLive(_10); // scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
79+
_10 = Option::<u16>::Some(const 99_u16); // scope 6 at $DIR/const_debuginfo.rs:+10:13: +10:24
80+
_17 = const 32_u32; // scope 7 at $DIR/const_debuginfo.rs:+12:13: +12:35
81+
_18 = const 32_u32; // scope 7 at $DIR/const_debuginfo.rs:+12:13: +12:35
82+
StorageLive(_11); // scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
83+
_11 = const 64_u32; // scope 8 at $DIR/const_debuginfo.rs:+13:13: +13:22
84+
StorageDead(_11); // scope 8 at $DIR/const_debuginfo.rs:+14:1: +14:2
85+
StorageDead(_10); // scope 6 at $DIR/const_debuginfo.rs:+14:1: +14:2
86+
StorageDead(_14); // scope 5 at $DIR/const_debuginfo.rs:+14:1: +14:2
87+
StorageDead(_15); // scope 5 at $DIR/const_debuginfo.rs:+14:1: +14:2
88+
StorageDead(_16); // scope 5 at $DIR/const_debuginfo.rs:+14:1: +14:2
8189
StorageDead(_9); // scope 4 at $DIR/const_debuginfo.rs:+14:1: +14:2
8290
StorageDead(_4); // scope 3 at $DIR/const_debuginfo.rs:+14:1: +14:2
8391
return; // scope 0 at $DIR/const_debuginfo.rs:+14:2: +14:2

0 commit comments

Comments
 (0)