use crate::transform::MirPass;
use rustc_index::bit_set::BitSet;
use rustc_index::vec::{Idx, IndexVec};
use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use smallvec::SmallVec;
use std::borrow::Cow;
use std::convert::TryInto;
pub struct SimplifyCfg {
label: String,
}
impl SimplifyCfg {
pub fn new(label: &str) -> Self {
SimplifyCfg { label: format!("SimplifyCfg-{}", label) }
}
}
pub fn simplify_cfg(body: &mut Body<'_>) {
CfgSimplifier::new(body).simplify();
remove_dead_blocks(body);
body.basic_blocks_mut().raw.shrink_to_fit();
}
impl<'tcx> MirPass<'tcx> for SimplifyCfg {
fn name(&self) -> Cow<'_, str> {
Cow::Borrowed(&self.label)
}
fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
debug!("SimplifyCfg({:?}) - simplifying {:?}", self.label, body.source);
simplify_cfg(body);
}
}
pub struct CfgSimplifier<'a, 'tcx> {
basic_blocks: &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
pred_count: IndexVec<BasicBlock, u32>,
}
impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
pub fn new(body: &'a mut Body<'tcx>) -> Self {
let mut pred_count = IndexVec::from_elem(0u32, body.basic_blocks());
pred_count[START_BLOCK] = 1;
for (_, data) in traversal::preorder(body) {
if let Some(ref term) = data.terminator {
for &tgt in term.successors() {
pred_count[tgt] += 1;
}
}
}
let basic_blocks = body.basic_blocks_mut();
CfgSimplifier { basic_blocks, pred_count }
}
pub fn simplify(mut self) {
self.strip_nops();
let mut start = START_BLOCK;
let mut merged_blocks = Vec::new();
loop {
let mut changed = false;
self.collapse_goto_chain(&mut start, &mut changed);
for bb in self.basic_blocks.indices() {
if self.pred_count[bb] == 0 {
continue;
}
debug!("simplifying {:?}", bb);
let mut terminator =
self.basic_blocks[bb].terminator.take().expect("invalid terminator state");
for successor in terminator.successors_mut() {
self.collapse_goto_chain(successor, &mut changed);
}
let mut inner_changed = true;
merged_blocks.clear();
while inner_changed {
inner_changed = false;
inner_changed |= self.simplify_branch(&mut terminator);
inner_changed |= self.merge_successor(&mut merged_blocks, &mut terminator);
changed |= inner_changed;
}
let statements_to_merge =
merged_blocks.iter().map(|&i| self.basic_blocks[i].statements.len()).sum();
if statements_to_merge > 0 {
let mut statements = std::mem::take(&mut self.basic_blocks[bb].statements);
statements.reserve(statements_to_merge);
for &from in &merged_blocks {
statements.append(&mut self.basic_blocks[from].statements);
}
self.basic_blocks[bb].statements = statements;
}
self.basic_blocks[bb].terminator = Some(terminator);
}
if !changed {
break;
}
}
if start != START_BLOCK {
debug_assert!(self.pred_count[START_BLOCK] == 0);
self.basic_blocks.swap(START_BLOCK, start);
self.pred_count.swap(START_BLOCK, start);
if self.pred_count[START_BLOCK] > 1 {
for (bb, data) in self.basic_blocks.iter_enumerated_mut() {
if self.pred_count[bb] == 0 {
continue;
}
for target in data.terminator_mut().successors_mut() {
if *target == start {
*target = START_BLOCK;
}
}
}
}
}
}
fn take_terminator_if_simple_goto(&mut self, bb: BasicBlock) -> Option<Terminator<'tcx>> {
match self.basic_blocks[bb] {
BasicBlockData {
ref statements,
terminator:
ref mut terminator @ Some(Terminator { kind: TerminatorKind::Goto { .. }, .. }),
..
} if statements.is_empty() => terminator.take(),
_ => None,
}
}
fn collapse_goto_chain(&mut self, start: &mut BasicBlock, changed: &mut bool) {
let mut terminators: SmallVec<[_; 1]> = Default::default();
let mut current = *start;
while let Some(terminator) = self.take_terminator_if_simple_goto(current) {
let target = match terminator {
Terminator { kind: TerminatorKind::Goto { target }, .. } => target,
_ => unreachable!(),
};
terminators.push((current, terminator));
current = target;
}
let last = current;
*start = last;
while let Some((current, mut terminator)) = terminators.pop() {
let target = match terminator {
Terminator { kind: TerminatorKind::Goto { ref mut target }, .. } => target,
_ => unreachable!(),
};
*changed |= *target != last;
*target = last;
debug!("collapsing goto chain from {:?} to {:?}", current, target);
if self.pred_count[current] == 1 {
self.pred_count[current] = 0;
} else {
self.pred_count[*target] += 1;
self.pred_count[current] -= 1;
}
self.basic_blocks[current].terminator = Some(terminator);
}
}
fn merge_successor(
&mut self,
merged_blocks: &mut Vec<BasicBlock>,
terminator: &mut Terminator<'tcx>,
) -> bool {
let target = match terminator.kind {
TerminatorKind::Goto { target } if self.pred_count[target] == 1 => target,
_ => return false,
};
debug!("merging block {:?} into {:?}", target, terminator);
*terminator = match self.basic_blocks[target].terminator.take() {
Some(terminator) => terminator,
None => {
return false;
}
};
merged_blocks.push(target);
self.pred_count[target] = 0;
true
}
fn simplify_branch(&mut self, terminator: &mut Terminator<'tcx>) -> bool {
match terminator.kind {
TerminatorKind::SwitchInt { .. } => {}
_ => return false,
};
let first_succ = {
if let Some(&first_succ) = terminator.successors().next() {
if terminator.successors().all(|s| *s == first_succ) {
let count = terminator.successors().count();
self.pred_count[first_succ] -= (count - 1) as u32;
first_succ
} else {
return false;
}
} else {
return false;
}
};
debug!("simplifying branch {:?}", terminator);
terminator.kind = TerminatorKind::Goto { target: first_succ };
true
}
fn strip_nops(&mut self) {
for blk in self.basic_blocks.iter_mut() {
blk.statements.retain(|stmt| !matches!(stmt.kind, StatementKind::Nop))
}
}
}
pub fn remove_dead_blocks(body: &mut Body<'_>) {
let mut seen = BitSet::new_empty(body.basic_blocks().len());
for (bb, _) in traversal::preorder(body) {
seen.insert(bb.index());
}
let basic_blocks = body.basic_blocks_mut();
let num_blocks = basic_blocks.len();
let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
let mut used_blocks = 0;
for alive_index in seen.iter() {
replacements[alive_index] = BasicBlock::new(used_blocks);
if alive_index != used_blocks {
basic_blocks.raw.swap(alive_index, used_blocks);
}
used_blocks += 1;
}
basic_blocks.raw.truncate(used_blocks);
for block in basic_blocks {
for target in block.terminator_mut().successors_mut() {
*target = replacements[target.index()];
}
}
}
pub struct SimplifyLocals;
impl<'tcx> MirPass<'tcx> for SimplifyLocals {
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
trace!("running SimplifyLocals on {:?}", body.source);
let mut used_locals = UsedLocals::new(body);
remove_unused_definitions(&mut used_locals, body);
let map = make_local_map(&mut body.local_decls, &used_locals);
if map.iter().any(Option::is_none) {
let mut updater = LocalUpdater { map, tcx };
updater.visit_body(body);
body.local_decls.shrink_to_fit();
}
}
}
fn make_local_map<V>(
local_decls: &mut IndexVec<Local, V>,
used_locals: &UsedLocals,
) -> IndexVec<Local, Option<Local>> {
let mut map: IndexVec<Local, Option<Local>> = IndexVec::from_elem(None, &*local_decls);
let mut used = Local::new(0);
for alive_index in local_decls.indices() {
if !used_locals.is_used(alive_index) {
continue;
}
map[alive_index] = Some(used);
if alive_index != used {
local_decls.swap(alive_index, used);
}
used.increment_by(1);
}
local_decls.truncate(used.index());
map
}
struct UsedLocals {
increment: bool,
arg_count: u32,
use_count: IndexVec<Local, u32>,
}
impl UsedLocals {
fn new(body: &Body<'_>) -> Self {
let mut this = Self {
increment: true,
arg_count: body.arg_count.try_into().unwrap(),
use_count: IndexVec::from_elem(0, &body.local_decls),
};
this.visit_body(body);
this
}
fn is_used(&self, local: Local) -> bool {
trace!("is_used({:?}): use_count: {:?}", local, self.use_count[local]);
local.as_u32() <= self.arg_count || self.use_count[local] != 0
}
fn statement_removed(&mut self, statement: &Statement<'tcx>) {
self.increment = false;
let location = Location { block: START_BLOCK, statement_index: 0 };
self.visit_statement(statement, location);
}
fn visit_lhs(&mut self, place: &Place<'tcx>, location: Location) {
if place.is_indirect() {
self.visit_place(place, PlaceContext::MutatingUse(MutatingUseContext::Store), location);
} else {
self.super_projection(
place.local,
&place.projection,
PlaceContext::MutatingUse(MutatingUseContext::Projection),
location,
);
}
}
}
impl Visitor<'_> for UsedLocals {
fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
match statement.kind {
StatementKind::LlvmInlineAsm(..)
| StatementKind::Retag(..)
| StatementKind::Coverage(..)
| StatementKind::FakeRead(..)
| StatementKind::AscribeUserType(..) => {
self.super_statement(statement, location);
}
StatementKind::Nop => {}
StatementKind::StorageLive(_local) | StatementKind::StorageDead(_local) => {}
StatementKind::Assign(box (ref place, ref rvalue)) => {
self.visit_lhs(place, location);
self.visit_rvalue(rvalue, location);
}
StatementKind::SetDiscriminant { ref place, variant_index: _ } => {
self.visit_lhs(place, location);
}
}
}
fn visit_local(&mut self, local: &Local, _ctx: PlaceContext, _location: Location) {
if self.increment {
self.use_count[*local] += 1;
} else {
assert_ne!(self.use_count[*local], 0);
self.use_count[*local] -= 1;
}
}
}
fn remove_unused_definitions<'a, 'tcx>(used_locals: &'a mut UsedLocals, body: &mut Body<'tcx>) {
let mut modified = true;
while modified {
modified = false;
for data in body.basic_blocks_mut() {
data.statements.retain(|statement| {
let keep = match &statement.kind {
StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => {
used_locals.is_used(*local)
}
StatementKind::Assign(box (place, _)) => used_locals.is_used(place.local),
StatementKind::SetDiscriminant { ref place, .. } => {
used_locals.is_used(place.local)
}
_ => true,
};
if !keep {
trace!("removing statement {:?}", statement);
modified = true;
used_locals.statement_removed(statement);
}
keep
});
}
}
}
struct LocalUpdater<'tcx> {
map: IndexVec<Local, Option<Local>>,
tcx: TyCtxt<'tcx>,
}
impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}
fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
*l = self.map[*l].unwrap();
}
}