use either::Either;
use rustc_data_structures::graph::dominators::Dominators;
use rustc_index::bit_set::BitSet;
use rustc_index::{IndexSlice, IndexVec};
use rustc_middle::middle::resolve_bound_vars::Set1;
use rustc_middle::mir::visit::*;
use rustc_middle::mir::*;
pub struct SsaLocals {
assignments: IndexVec<Local, Set1<LocationExtended>>,
assignment_order: Vec<Local>,
copy_classes: IndexVec<Local, Local>,
direct_uses: IndexVec<Local, u32>,
}
struct SmallDominators<'a> {
inner: Option<&'a Dominators<BasicBlock>>,
}
impl SmallDominators<'_> {
fn dominates(&self, first: Location, second: Location) -> bool {
if first.block == second.block {
first.statement_index <= second.statement_index
} else if let Some(inner) = &self.inner {
inner.dominates(first.block, second.block)
} else {
first.block < second.block
}
}
fn check_dominates(&mut self, set: &mut Set1<LocationExtended>, loc: Location) {
let assign_dominates = match *set {
Set1::Empty | Set1::Many => false,
Set1::One(LocationExtended::Arg) => true,
Set1::One(LocationExtended::Plain(assign)) => {
self.dominates(assign.successor_within_block(), loc)
}
};
if !assign_dominates {
*set = Set1::Many;
}
}
}
impl SsaLocals {
pub fn new<'tcx>(body: &Body<'tcx>) -> SsaLocals {
let assignment_order = Vec::with_capacity(body.local_decls.len());
let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
let dominators =
if body.basic_blocks.len() > 2 { Some(body.basic_blocks.dominators()) } else { None };
let dominators = SmallDominators { inner: dominators };
let direct_uses = IndexVec::from_elem(0, &body.local_decls);
let mut visitor = SsaVisitor { assignments, assignment_order, dominators, direct_uses };
for local in body.args_iter() {
visitor.assignments[local] = Set1::One(LocationExtended::Arg);
}
for (bb, data) in traversal::reverse_postorder(body) {
visitor.visit_basic_block_data(bb, data);
}
for var_debug_info in &body.var_debug_info {
visitor.visit_var_debug_info(var_debug_info);
}
debug!(?visitor.assignments);
debug!(?visitor.direct_uses);
visitor
.assignment_order
.retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
debug!(?visitor.assignment_order);
let mut ssa = SsaLocals {
assignments: visitor.assignments,
assignment_order: visitor.assignment_order,
direct_uses: visitor.direct_uses,
copy_classes: IndexVec::default(),
};
compute_copy_classes(&mut ssa, body);
ssa
}
pub fn num_locals(&self) -> usize {
self.assignments.len()
}
pub fn locals(&self) -> impl Iterator<Item = Local> {
self.assignments.indices()
}
pub fn is_ssa(&self, local: Local) -> bool {
matches!(self.assignments[local], Set1::One(_))
}
pub fn num_direct_uses(&self, local: Local) -> u32 {
self.direct_uses[local]
}
pub fn assignment_dominates(
&self,
dominators: &Dominators<BasicBlock>,
local: Local,
location: Location,
) -> bool {
match self.assignments[local] {
Set1::One(LocationExtended::Arg) => true,
Set1::One(LocationExtended::Plain(ass)) => {
if ass.block == location.block {
ass.statement_index < location.statement_index
} else {
dominators.dominates(ass.block, location.block)
}
}
_ => false,
}
}
pub fn assignments<'a, 'tcx>(
&'a self,
body: &'a Body<'tcx>,
) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a {
self.assignment_order.iter().filter_map(|&local| {
if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
assert_eq!(target.as_local(), Some(local));
Some((local, rvalue, loc))
} else {
None
}
})
}
pub fn for_each_assignment_mut<'tcx>(
&self,
basic_blocks: &mut BasicBlocks<'tcx>,
mut f: impl FnMut(Local, &mut Rvalue<'tcx>, Location),
) {
for &local in &self.assignment_order {
if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
let bbs = basic_blocks.as_mut_preserves_cfg();
let bb = &mut bbs[loc.block];
let stmt = &mut bb.statements[loc.statement_index];
let StatementKind::Assign(box (target, ref mut rvalue)) = stmt.kind else { bug!() };
assert_eq!(target.as_local(), Some(local));
f(local, rvalue, loc)
}
}
}
pub fn copy_classes(&self) -> &IndexSlice<Local, Local> {
&self.copy_classes
}
pub fn meet_copy_equivalence(&self, property: &mut BitSet<Local>) {
for (local, &head) in self.copy_classes.iter_enumerated() {
if !property.contains(local) {
property.remove(head);
}
}
for (local, &head) in self.copy_classes.iter_enumerated() {
if !property.contains(head) {
property.remove(local);
}
}
#[cfg(debug_assertions)]
for (local, &head) in self.copy_classes.iter_enumerated() {
assert_eq!(property.contains(local), property.contains(head));
}
}
}
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum LocationExtended {
Plain(Location),
Arg,
}
struct SsaVisitor<'a> {
dominators: SmallDominators<'a>,
assignments: IndexVec<Local, Set1<LocationExtended>>,
assignment_order: Vec<Local>,
direct_uses: IndexVec<Local, u32>,
}
impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
match ctxt {
PlaceContext::MutatingUse(MutatingUseContext::Projection)
| PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection) => bug!(),
PlaceContext::NonMutatingUse(
NonMutatingUseContext::SharedBorrow
| NonMutatingUseContext::ShallowBorrow
| NonMutatingUseContext::AddressOf,
)
| PlaceContext::MutatingUse(_) => {
self.assignments[local] = Set1::Many;
}
PlaceContext::NonMutatingUse(_) => {
self.dominators.check_dominates(&mut self.assignments[local], loc);
self.direct_uses[local] += 1;
}
PlaceContext::NonUse(_) => {}
}
}
fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) {
if place.projection.first() == Some(&PlaceElem::Deref) {
if ctxt.is_use() {
let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
self.visit_projection(place.as_ref(), new_ctxt, loc);
self.dominators.check_dominates(&mut self.assignments[place.local], loc);
}
return;
} else {
self.visit_projection(place.as_ref(), ctxt, loc);
self.visit_local(place.local, ctxt, loc);
}
}
fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, loc: Location) {
if let Some(local) = place.as_local() {
self.assignments[local].insert(LocationExtended::Plain(loc));
if let Set1::One(_) = self.assignments[local] {
self.assignment_order.push(local);
}
} else {
self.visit_place(place, PlaceContext::MutatingUse(MutatingUseContext::Store), loc);
}
self.visit_rvalue(rvalue, loc);
}
}
#[instrument(level = "trace", skip(ssa, body))]
fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
let mut direct_uses = std::mem::take(&mut ssa.direct_uses);
let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
for (local, rvalue, _) in ssa.assignments(body) {
let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
| Rvalue::CopyForDeref(place)) = rvalue
else {
continue;
};
let Some(rhs) = place.as_local() else { continue };
let local_ty = body.local_decls()[local].ty;
let rhs_ty = body.local_decls()[rhs].ty;
if local_ty != rhs_ty {
trace!("skipped `{local:?} = {rhs:?}` due to subtyping: {local_ty} != {rhs_ty}");
continue;
}
if !ssa.is_ssa(rhs) {
continue;
}
let head = copies[rhs];
if local == RETURN_PLACE {
if body.local_kind(head) != LocalKind::Temp {
continue;
}
for h in copies.iter_mut() {
if *h == head {
*h = RETURN_PLACE;
}
}
} else {
copies[local] = head;
}
direct_uses[rhs] -= 1;
}
debug!(?copies);
debug!(?direct_uses);
#[cfg(debug_assertions)]
for &head in copies.iter() {
assert_eq!(copies[head], head);
}
debug_assert_eq!(copies[RETURN_PLACE], RETURN_PLACE);
ssa.direct_uses = direct_uses;
ssa.copy_classes = copies;
}
#[derive(Debug)]
pub(crate) struct StorageLiveLocals {
storage_live: IndexVec<Local, Set1<LocationExtended>>,
}
impl StorageLiveLocals {
pub(crate) fn new(
body: &Body<'_>,
always_storage_live_locals: &BitSet<Local>,
) -> StorageLiveLocals {
let mut storage_live = IndexVec::from_elem(Set1::Empty, &body.local_decls);
for local in always_storage_live_locals.iter() {
storage_live[local] = Set1::One(LocationExtended::Arg);
}
for (block, bbdata) in body.basic_blocks.iter_enumerated() {
for (statement_index, statement) in bbdata.statements.iter().enumerate() {
if let StatementKind::StorageLive(local) = statement.kind {
storage_live[local]
.insert(LocationExtended::Plain(Location { block, statement_index }));
}
}
}
debug!(?storage_live);
StorageLiveLocals { storage_live }
}
#[inline]
pub(crate) fn has_single_storage(&self, local: Local) -> bool {
matches!(self.storage_live[local], Set1::One(_))
}
}