Skip to content

Commit 0843acb

Browse files
cjgillotcjgillot
committed
Fix SROA without deaggregation.
1 parent 3de7d7f commit 0843acb

10 files changed

+174
-127
lines changed

‎compiler/rustc_mir_transform/src/sroa.rs

+82-45
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use crate::MirPass;
22
use rustc_data_structures::fx::{FxIndexMap, IndexEntry};
33
use rustc_index::bit_set::BitSet;
44
use rustc_index::vec::IndexVec;
5+
use rustc_middle::mir::patch::MirPatch;
56
use rustc_middle::mir::visit::*;
67
use rustc_middle::mir::*;
78
use rustc_middle::ty::TyCtxt;
@@ -13,7 +14,9 @@ impl<'tcx> MirPass<'tcx> for ScalarReplacementOfAggregates {
1314
sess.mir_opt_level() >= 3
1415
}
1516

17+
#[instrument(level = "debug", skip(self, tcx, body))]
1618
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
19+
debug!(def_id = ?body.source.def_id());
1720
let escaping = escaping_locals(&*body);
1821
debug!(?escaping);
1922
let replacements = compute_flattening(tcx, body, escaping);
@@ -69,15 +72,28 @@ fn escaping_locals(body: &Body<'_>) -> BitSet<Local> {
6972
self.super_rvalue(rvalue, location)
7073
}
7174

75+
fn visit_assign(
76+
&mut self,
77+
lvalue: &Place<'tcx>,
78+
rvalue: &Rvalue<'tcx>,
79+
location: Location,
80+
) {
81+
if lvalue.as_local().is_some() && let Rvalue::Aggregate(..) = rvalue {
82+
// Aggregate assignments are expanded in run_pass.
83+
self.visit_rvalue(rvalue, location);
84+
return;
85+
}
86+
self.super_assign(lvalue, rvalue, location)
87+
}
88+
7289
fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
73-
if let StatementKind::StorageLive(..)
74-
| StatementKind::StorageDead(..)
75-
| StatementKind::Deinit(..) = statement.kind
76-
{
90+
match statement.kind {
7791
// Storage statements are expanded in run_pass.
78-
return;
92+
StatementKind::StorageLive(..)
93+
| StatementKind::StorageDead(..)
94+
| StatementKind::Deinit(..) => return,
95+
_ => self.super_statement(statement, location),
7996
}
80-
self.super_statement(statement, location)
8197
}
8298

8399
fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
@@ -192,6 +208,7 @@ fn replace_flattened_locals<'tcx>(
192208
replacements,
193209
all_dead_locals,
194210
fragments,
211+
patch: MirPatch::new(body),
195212
};
196213
for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
197214
visitor.visit_basic_block_data(bb, data);
@@ -205,6 +222,7 @@ fn replace_flattened_locals<'tcx>(
205222
for var_debug_info in &mut body.var_debug_info {
206223
visitor.visit_var_debug_info(var_debug_info);
207224
}
225+
visitor.patch.apply(body);
208226
}
209227

210228
struct ReplacementVisitor<'tcx, 'll> {
@@ -218,6 +236,7 @@ struct ReplacementVisitor<'tcx, 'll> {
218236
/// Pre-computed list of all "new" locals for each "old" local. This is used to expand storage
219237
/// and deinit statement and debuginfo.
220238
fragments: IndexVec<Local, Vec<(&'tcx [PlaceElem<'tcx>], Local)>>,
239+
patch: MirPatch<'tcx>,
221240
}
222241

223242
impl<'tcx, 'll> ReplacementVisitor<'tcx, 'll> {
@@ -255,12 +274,63 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
255274
}
256275

257276
fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
258-
if let StatementKind::StorageLive(..)
259-
| StatementKind::StorageDead(..)
260-
| StatementKind::Deinit(..) = statement.kind
261-
{
262-
// Storage statements are expanded in run_pass.
263-
return;
277+
match statement.kind {
278+
StatementKind::StorageLive(l) => {
279+
if self.all_dead_locals.contains(l) {
280+
let final_locals = &self.fragments[l];
281+
for &(_, fl) in final_locals {
282+
self.patch.add_statement(location, StatementKind::StorageLive(fl));
283+
}
284+
statement.make_nop();
285+
}
286+
return;
287+
}
288+
StatementKind::StorageDead(l) => {
289+
if self.all_dead_locals.contains(l) {
290+
let final_locals = &self.fragments[l];
291+
for &(_, fl) in final_locals {
292+
self.patch.add_statement(location, StatementKind::StorageDead(fl));
293+
}
294+
statement.make_nop();
295+
}
296+
return;
297+
}
298+
StatementKind::Deinit(box ref place) => {
299+
if let Some(local) = place.as_local()
300+
&& self.all_dead_locals.contains(local)
301+
{
302+
let final_locals = &self.fragments[local];
303+
for &(_, fl) in final_locals {
304+
self.patch.add_statement(
305+
location,
306+
StatementKind::Deinit(Box::new(fl.into())),
307+
);
308+
}
309+
statement.make_nop();
310+
return;
311+
}
312+
}
313+
314+
StatementKind::Assign(box (ref place, Rvalue::Aggregate(_, ref operands))) => {
315+
if let Some(local) = place.as_local()
316+
&& self.all_dead_locals.contains(local)
317+
{
318+
let final_locals = &self.fragments[local];
319+
for &(projection, fl) in final_locals {
320+
let &[PlaceElem::Field(index, _)] = projection else { bug!() };
321+
let index = index.as_usize();
322+
let rvalue = Rvalue::Use(operands[index].clone());
323+
self.patch.add_statement(
324+
location,
325+
StatementKind::Assign(Box::new((fl.into(), rvalue))),
326+
);
327+
}
328+
statement.make_nop();
329+
return;
330+
}
331+
}
332+
333+
_ => {}
264334
}
265335
self.super_statement(statement, location)
266336
}
@@ -309,39 +379,6 @@ impl<'tcx, 'll> MutVisitor<'tcx> for ReplacementVisitor<'tcx, 'll> {
309379
}
310380
}
311381

312-
fn visit_basic_block_data(&mut self, bb: BasicBlock, bbdata: &mut BasicBlockData<'tcx>) {
313-
self.super_basic_block_data(bb, bbdata);
314-
315-
#[derive(Debug)]
316-
enum Stmt {
317-
StorageLive,
318-
StorageDead,
319-
Deinit,
320-
}
321-
322-
bbdata.expand_statements(|stmt| {
323-
let source_info = stmt.source_info;
324-
let (stmt, origin_local) = match &stmt.kind {
325-
StatementKind::StorageLive(l) => (Stmt::StorageLive, *l),
326-
StatementKind::StorageDead(l) => (Stmt::StorageDead, *l),
327-
StatementKind::Deinit(p) if let Some(l) = p.as_local() => (Stmt::Deinit, l),
328-
_ => return None,
329-
};
330-
if !self.all_dead_locals.contains(origin_local) {
331-
return None;
332-
}
333-
let final_locals = self.fragments.get(origin_local)?;
334-
Some(final_locals.iter().map(move |&(_, l)| {
335-
let kind = match stmt {
336-
Stmt::StorageLive => StatementKind::StorageLive(l),
337-
Stmt::StorageDead => StatementKind::StorageDead(l),
338-
Stmt::Deinit => StatementKind::Deinit(Box::new(l.into())),
339-
};
340-
Statement { source_info, kind }
341-
}))
342-
});
343-
}
344-
345382
fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
346383
assert!(!self.all_dead_locals.contains(*local));
347384
}

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

+12-18
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 _14: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:13: +13:16
12-
let mut _15: u32; // in scope 0 at $DIR/const_debuginfo.rs:+13:19: +13:22
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
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
@@ -35,12 +35,13 @@
3535
let _11: std::option::Option<u16>; // in scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
3636
scope 7 {
3737
debug o => _11; // in scope 7 at $DIR/const_debuginfo.rs:+10:9: +10:10
38-
let _12: Point; // in scope 7 at $DIR/const_debuginfo.rs:+12:9: +12: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
3940
scope 8 {
40-
debug p => _12; // in scope 8 at $DIR/const_debuginfo.rs:+12:9: +12:10
41-
let _13: u32; // in scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
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
4243
scope 9 {
43-
- debug a => _13; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
44+
- debug a => _12; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
4445
+ debug a => const 64_u32; // in scope 9 at $DIR/const_debuginfo.rs:+13:9: +13:10
4546
}
4647
}
@@ -70,18 +71,11 @@
7071
_10 = (const true, const false, const 123_u32); // scope 5 at $DIR/const_debuginfo.rs:+8:13: +8:34
7172
StorageLive(_11); // scope 6 at $DIR/const_debuginfo.rs:+10:9: +10:10
7273
_11 = Option::<u16>::Some(const 99_u16); // scope 6 at $DIR/const_debuginfo.rs:+10:13: +10:24
73-
StorageLive(_12); // scope 7 at $DIR/const_debuginfo.rs:+12:9: +12:10
74-
_12 = Point { x: const 32_u32, y: const 32_u32 }; // scope 7 at $DIR/const_debuginfo.rs:+12:13: +12:35
75-
StorageLive(_13); // scope 8 at $DIR/const_debuginfo.rs:+13:9: +13:10
76-
StorageLive(_14); // scope 8 at $DIR/const_debuginfo.rs:+13:13: +13:16
77-
_14 = const 32_u32; // scope 8 at $DIR/const_debuginfo.rs:+13:13: +13:16
78-
StorageLive(_15); // scope 8 at $DIR/const_debuginfo.rs:+13:19: +13:22
79-
_15 = const 32_u32; // scope 8 at $DIR/const_debuginfo.rs:+13:19: +13:22
80-
_13 = const 64_u32; // scope 8 at $DIR/const_debuginfo.rs:+13:13: +13:22
81-
StorageDead(_15); // scope 8 at $DIR/const_debuginfo.rs:+13:21: +13:22
82-
StorageDead(_14); // scope 8 at $DIR/const_debuginfo.rs:+13:21: +13:22
83-
StorageDead(_13); // scope 8 at $DIR/const_debuginfo.rs:+14:1: +14:2
84-
StorageDead(_12); // scope 7 at $DIR/const_debuginfo.rs:+14:1: +14:2
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
8579
StorageDead(_11); // scope 6 at $DIR/const_debuginfo.rs:+14:1: +14:2
8680
StorageDead(_10); // scope 5 at $DIR/const_debuginfo.rs:+14:1: +14:2
8781
StorageDead(_9); // scope 4 at $DIR/const_debuginfo.rs:+14:1: +14:2

‎tests/mir-opt/const_prop/mutable_variable_unprop_assign.main.ConstProp.diff

+19-22
Original file line numberDiff line numberDiff line change
@@ -4,18 +4,19 @@
44
fn main() -> () {
55
let mut _0: (); // return place in scope 0 at $DIR/mutable_variable_unprop_assign.rs:+0:11: +0:11
66
let _1: i32; // in scope 0 at $DIR/mutable_variable_unprop_assign.rs:+1:9: +1:10
7-
let mut _3: i32; // in scope 0 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
7+
let mut _2: i32; // in scope 0 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
88
scope 1 {
99
debug a => _1; // in scope 1 at $DIR/mutable_variable_unprop_assign.rs:+1:9: +1:10
10-
let mut _2: (i32, i32); // in scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
10+
let mut _5: i32; // in scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
11+
let mut _6: i32; // in scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
1112
scope 2 {
12-
debug x => _2; // in scope 2 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
13-
let _4: i32; // in scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
13+
debug x => (i32, i32){ .1 => _5, .0 => _6, }; // in scope 2 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
14+
let _3: i32; // in scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
1415
scope 3 {
15-
debug y => _4; // in scope 3 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
16-
let _5: i32; // in scope 3 at $DIR/mutable_variable_unprop_assign.rs:+5:9: +5:10
16+
debug y => _3; // in scope 3 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
17+
let _4: i32; // in scope 3 at $DIR/mutable_variable_unprop_assign.rs:+5:9: +5:10
1718
scope 4 {
18-
debug z => _5; // in scope 4 at $DIR/mutable_variable_unprop_assign.rs:+5:9: +5:10
19+
debug z => _6; // in scope 4 at $DIR/mutable_variable_unprop_assign.rs:+5:9: +5:10
1920
}
2021
}
2122
}
@@ -30,21 +31,17 @@
3031
}
3132

3233
bb1: {
33-
StorageLive(_2); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
34-
- _2 = (const 1_i32, const 2_i32); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:29: +2:35
35-
+ _2 = const (1_i32, 2_i32); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:29: +2:35
36-
StorageLive(_3); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
37-
_3 = _1; // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
38-
(_2.1: i32) = move _3; // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:5: +3:12
39-
StorageDead(_3); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
40-
StorageLive(_4); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
41-
_4 = (_2.1: i32); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:13: +4:16
42-
StorageLive(_5); // scope 3 at $DIR/mutable_variable_unprop_assign.rs:+5:9: +5:10
43-
- _5 = (_2.0: i32); // scope 3 at $DIR/mutable_variable_unprop_assign.rs:+5:13: +5:16
44-
+ _5 = const 1_i32; // scope 3 at $DIR/mutable_variable_unprop_assign.rs:+5:13: +5:16
45-
StorageDead(_5); // scope 3 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
46-
StorageDead(_4); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
47-
StorageDead(_2); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
34+
StorageLive(_5); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:9: +2:14
35+
_5 = const 2_i32; // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:29: +2:35
36+
_6 = const 1_i32; // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+2:29: +2:35
37+
StorageLive(_2); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
38+
_2 = _1; // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
39+
_5 = move _2; // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:5: +3:12
40+
StorageDead(_2); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+3:11: +3:12
41+
StorageLive(_3); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:9: +4:10
42+
_3 = _5; // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+4:13: +4:16
43+
StorageDead(_3); // scope 2 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
44+
StorageDead(_5); // scope 1 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
4845
StorageDead(_1); // scope 0 at $DIR/mutable_variable_unprop_assign.rs:+6:1: +6:2
4946
return; // scope 0 at $DIR/mutable_variable_unprop_assign.rs:+6:2: +6:2
5047
}

‎tests/mir-opt/const_prop/optimizes_into_variable.main.ConstProp.64bit.diff

+3-9
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,15 @@
99
let _5: usize; // in scope 0 at $DIR/optimizes_into_variable.rs:+2:32: +2:33
1010
let mut _6: usize; // in scope 0 at $DIR/optimizes_into_variable.rs:+2:13: +2:34
1111
let mut _7: bool; // in scope 0 at $DIR/optimizes_into_variable.rs:+2:13: +2:34
12-
let mut _9: Point; // in scope 0 at $DIR/optimizes_into_variable.rs:+3:13: +3:36
12+
let mut _9: u32; // in scope 0 at $DIR/optimizes_into_variable.rs:+3:13: +3:36
1313
scope 1 {
1414
debug x => _1; // in scope 1 at $DIR/optimizes_into_variable.rs:+1:9: +1:10
1515
let _3: i32; // in scope 1 at $DIR/optimizes_into_variable.rs:+2:9: +2:10
1616
scope 2 {
1717
debug y => _3; // in scope 2 at $DIR/optimizes_into_variable.rs:+2:9: +2:10
1818
let _8: u32; // in scope 2 at $DIR/optimizes_into_variable.rs:+3:9: +3:10
1919
scope 3 {
20-
debug z => _8; // in scope 3 at $DIR/optimizes_into_variable.rs:+3:9: +3:10
20+
debug z => _9; // in scope 3 at $DIR/optimizes_into_variable.rs:+3:9: +3:10
2121
}
2222
}
2323
}
@@ -50,13 +50,7 @@
5050
+ _3 = const 3_i32; // scope 1 at $DIR/optimizes_into_variable.rs:+2:13: +2:34
5151
StorageDead(_5); // scope 1 at $DIR/optimizes_into_variable.rs:+2:34: +2:35
5252
StorageDead(_4); // scope 1 at $DIR/optimizes_into_variable.rs:+2:34: +2:35
53-
StorageLive(_8); // scope 2 at $DIR/optimizes_into_variable.rs:+3:9: +3:10
54-
StorageLive(_9); // scope 2 at $DIR/optimizes_into_variable.rs:+3:13: +3:36
55-
_9 = Point { x: const 12_u32, y: const 42_u32 }; // scope 2 at $DIR/optimizes_into_variable.rs:+3:13: +3:36
56-
- _8 = (_9.1: u32); // scope 2 at $DIR/optimizes_into_variable.rs:+3:13: +3:38
57-
+ _8 = const 42_u32; // scope 2 at $DIR/optimizes_into_variable.rs:+3:13: +3:38
58-
StorageDead(_9); // scope 2 at $DIR/optimizes_into_variable.rs:+3:38: +3:39
59-
StorageDead(_8); // scope 2 at $DIR/optimizes_into_variable.rs:+4:1: +4:2
53+
_9 = const 42_u32; // scope 2 at $DIR/optimizes_into_variable.rs:+3:13: +3:36
6054
StorageDead(_3); // scope 1 at $DIR/optimizes_into_variable.rs:+4:1: +4:2
6155
StorageDead(_1); // scope 0 at $DIR/optimizes_into_variable.rs:+4:1: +4:2
6256
return; // scope 0 at $DIR/optimizes_into_variable.rs:+4:2: +4:2

0 commit comments

Comments
 (0)