Skip to content

Commit dec8b9d

Browse files
dianqkdianqk
committed
Updating the MIR with MirPatch
1 parent f79fd40 commit dec8b9d

File tree

1 file changed

+52
-50
lines changed

1 file changed

+52
-50
lines changed

‎compiler/rustc_mir_transform/src/match_branches.rs

+52-50
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
use rustc_index::IndexVec;
1+
use rustc_index::IndexSlice;
2+
use rustc_middle::mir::patch::MirPatch;
23
use rustc_middle::mir::*;
34
use rustc_middle::ty::{ParamEnv, ScalarInt, Ty, TyCtxt};
45
use std::iter;
@@ -16,9 +17,10 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
1617
let def_id = body.source.def_id();
1718
let param_env = tcx.param_env_reveal_all_normalized(def_id);
1819

19-
let bbs = body.basic_blocks.as_mut();
2020
let mut should_cleanup = false;
21-
for bb_idx in bbs.indices() {
21+
for i in 0..body.basic_blocks.len() {
22+
let bbs = &*body.basic_blocks;
23+
let bb_idx = BasicBlock::from_usize(i);
2224
if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {def_id:?} ")) {
2325
continue;
2426
}
@@ -34,12 +36,11 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
3436
_ => continue,
3537
};
3638

37-
if SimplifyToIf.simplify(tcx, &mut body.local_decls, bbs, bb_idx, param_env) {
39+
if SimplifyToIf.simplify(tcx, body, bb_idx, param_env) {
3840
should_cleanup = true;
3941
continue;
4042
}
41-
if SimplifyToExp::default().simplify(tcx, &mut body.local_decls, bbs, bb_idx, param_env)
42-
{
43+
if SimplifyToExp::default().simplify(tcx, body, bb_idx, param_env) {
4344
should_cleanup = true;
4445
continue;
4546
}
@@ -57,43 +58,42 @@ trait SimplifyMatch<'tcx> {
5758
fn simplify(
5859
&mut self,
5960
tcx: TyCtxt<'tcx>,
60-
local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
61-
bbs: &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
61+
body: &mut Body<'tcx>,
6262
switch_bb_idx: BasicBlock,
6363
param_env: ParamEnv<'tcx>,
6464
) -> bool {
65+
let bbs = &body.basic_blocks;
6566
let (discr, targets) = match bbs[switch_bb_idx].terminator().kind {
6667
TerminatorKind::SwitchInt { ref discr, ref targets, .. } => (discr, targets),
6768
_ => unreachable!(),
6869
};
6970

70-
let discr_ty = discr.ty(local_decls, tcx);
71+
// let local_decls = body.local_decls();
72+
let discr_ty = discr.ty(body.local_decls(), tcx);
7173
if !self.can_simplify(tcx, targets, param_env, bbs, discr_ty) {
7274
return false;
7375
}
7476

77+
let mut patch = MirPatch::new(body);
78+
7579
// Take ownership of items now that we know we can optimize.
7680
let discr = discr.clone();
7781

7882
// Introduce a temporary for the discriminant value.
7983
let source_info = bbs[switch_bb_idx].terminator().source_info;
80-
let discr_local = local_decls.push(LocalDecl::new(discr_ty, source_info.span));
84+
let discr_local = patch.new_temp(discr_ty, source_info.span);
8185

8286
// We already checked that targets are different blocks,
8387
// and bb_idx has a different terminator from both of them.
84-
let new_stmts = self.new_stmts(tcx, targets, param_env, bbs, discr_local, discr_ty);
8588
let (_, first) = targets.iter().next().unwrap();
86-
let (from, first) = bbs.pick2_mut(switch_bb_idx, first);
87-
from.statements
88-
.push(Statement { source_info, kind: StatementKind::StorageLive(discr_local) });
89-
from.statements.push(Statement {
90-
source_info,
91-
kind: StatementKind::Assign(Box::new((Place::from(discr_local), Rvalue::Use(discr)))),
92-
});
93-
from.statements.extend(new_stmts);
94-
from.statements
95-
.push(Statement { source_info, kind: StatementKind::StorageDead(discr_local) });
96-
from.terminator_mut().kind = first.terminator().kind.clone();
89+
let statement_index = bbs[switch_bb_idx].statements.len();
90+
let parent_end = Location { block: switch_bb_idx, statement_index };
91+
patch.add_statement(parent_end, StatementKind::StorageLive(discr_local));
92+
patch.add_assign(parent_end, Place::from(discr_local), Rvalue::Use(discr));
93+
self.new_stmts(tcx, targets, param_env, &mut patch, parent_end, bbs, discr_local, discr_ty);
94+
patch.add_statement(parent_end, StatementKind::StorageDead(discr_local));
95+
patch.patch_terminator(switch_bb_idx, bbs[first].terminator().kind.clone());
96+
patch.apply(body);
9797
true
9898
}
9999

@@ -102,7 +102,7 @@ trait SimplifyMatch<'tcx> {
102102
tcx: TyCtxt<'tcx>,
103103
targets: &SwitchTargets,
104104
param_env: ParamEnv<'tcx>,
105-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
105+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
106106
discr_ty: Ty<'tcx>,
107107
) -> bool;
108108

@@ -111,10 +111,12 @@ trait SimplifyMatch<'tcx> {
111111
tcx: TyCtxt<'tcx>,
112112
targets: &SwitchTargets,
113113
param_env: ParamEnv<'tcx>,
114-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
114+
patch: &mut MirPatch<'tcx>,
115+
parent_end: Location,
116+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
115117
discr_local: Local,
116118
discr_ty: Ty<'tcx>,
117-
) -> Vec<Statement<'tcx>>;
119+
);
118120
}
119121

120122
struct SimplifyToIf;
@@ -156,7 +158,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
156158
tcx: TyCtxt<'tcx>,
157159
targets: &SwitchTargets,
158160
param_env: ParamEnv<'tcx>,
159-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
161+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
160162
_discr_ty: Ty<'tcx>,
161163
) -> bool {
162164
if targets.iter().len() != 1 {
@@ -207,20 +209,23 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
207209
tcx: TyCtxt<'tcx>,
208210
targets: &SwitchTargets,
209211
param_env: ParamEnv<'tcx>,
210-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
212+
patch: &mut MirPatch<'tcx>,
213+
parent_end: Location,
214+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
211215
discr_local: Local,
212216
discr_ty: Ty<'tcx>,
213-
) -> Vec<Statement<'tcx>> {
217+
) {
214218
let (val, first) = targets.iter().next().unwrap();
215219
let second = targets.otherwise();
216220
// We already checked that first and second are different blocks,
217221
// and bb_idx has a different terminator from both of them.
218222
let first = &bbs[first];
219223
let second = &bbs[second];
220-
221-
let new_stmts = iter::zip(&first.statements, &second.statements).map(|(f, s)| {
224+
for (f, s) in iter::zip(&first.statements, &second.statements) {
222225
match (&f.kind, &s.kind) {
223-
(f_s, s_s) if f_s == s_s => (*f).clone(),
226+
(f_s, s_s) if f_s == s_s => {
227+
patch.add_statement(parent_end, f.kind.clone());
228+
}
224229

225230
(
226231
StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
@@ -231,7 +236,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
231236
let s_b = s_c.const_.try_eval_bool(tcx, param_env).unwrap();
232237
if f_b == s_b {
233238
// Same value in both blocks. Use statement as is.
234-
(*f).clone()
239+
patch.add_statement(parent_end, f.kind.clone());
235240
} else {
236241
// Different value between blocks. Make value conditional on switch condition.
237242
let size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size;
@@ -246,17 +251,13 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
246251
op,
247252
Box::new((Operand::Copy(Place::from(discr_local)), const_cmp)),
248253
);
249-
Statement {
250-
source_info: f.source_info,
251-
kind: StatementKind::Assign(Box::new((*lhs, rhs))),
252-
}
254+
patch.add_assign(parent_end, *lhs, rhs);
253255
}
254256
}
255257

256258
_ => unreachable!(),
257259
}
258-
});
259-
new_stmts.collect()
260+
}
260261
}
261262
}
262263

@@ -333,7 +334,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
333334
tcx: TyCtxt<'tcx>,
334335
targets: &SwitchTargets,
335336
param_env: ParamEnv<'tcx>,
336-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
337+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
337338
discr_ty: Ty<'tcx>,
338339
) -> bool {
339340
if targets.iter().len() < 2 || targets.iter().len() > 64 {
@@ -467,16 +468,20 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
467468
_tcx: TyCtxt<'tcx>,
468469
targets: &SwitchTargets,
469470
_param_env: ParamEnv<'tcx>,
470-
bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
471+
patch: &mut MirPatch<'tcx>,
472+
parent_end: Location,
473+
bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
471474
discr_local: Local,
472475
discr_ty: Ty<'tcx>,
473-
) -> Vec<Statement<'tcx>> {
476+
) {
474477
let (_, first) = targets.iter().next().unwrap();
475478
let first = &bbs[first];
476479

477-
let new_stmts =
478-
iter::zip(&self.transfrom_types, &first.statements).map(|(t, s)| match (t, &s.kind) {
479-
(TransfromType::Same, _) | (TransfromType::Eq, _) => (*s).clone(),
480+
for (t, s) in iter::zip(&self.transfrom_types, &first.statements) {
481+
match (t, &s.kind) {
482+
(TransfromType::Same, _) | (TransfromType::Eq, _) => {
483+
patch.add_statement(parent_end, s.kind.clone());
484+
}
480485
(
481486
TransfromType::Discr,
482487
StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
@@ -487,13 +492,10 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
487492
} else {
488493
Rvalue::Cast(CastKind::IntToInt, operand, f_c.const_.ty())
489494
};
490-
Statement {
491-
source_info: s.source_info,
492-
kind: StatementKind::Assign(Box::new((*lhs, r_val))),
493-
}
495+
patch.add_assign(parent_end, *lhs, r_val);
494496
}
495497
_ => unreachable!(),
496-
});
497-
new_stmts.collect()
498+
}
499+
}
498500
}
499501
}

0 commit comments

Comments
 (0)