use rustc_index::IndexVec;
use rustc_middle::mir::visit::{PlaceContext, Visitor};
use rustc_middle::mir::{Body, Local, Location};
use rustc_mir_dataflow::points::{DenseLocationMap, PointIndex};
use crate::def_use::{self, DefUse};
pub(crate) struct LocalUseMap {
first_def_at: IndexVec<Local, Option<AppearanceIndex>>,
first_use_at: IndexVec<Local, Option<AppearanceIndex>>,
first_drop_at: IndexVec<Local, Option<AppearanceIndex>>,
appearances: Appearances,
}
type Appearances = IndexVec<AppearanceIndex, Appearance>;
struct Appearance {
point_index: PointIndex,
next: Option<AppearanceIndex>,
}
rustc_index::newtype_index! {
pub struct AppearanceIndex {}
}
fn appearances_iter(
first: Option<AppearanceIndex>,
appearances: &Appearances,
) -> impl Iterator<Item = AppearanceIndex> + '_ {
AppearancesIter { appearances, current: first }
}
struct AppearancesIter<'a> {
appearances: &'a Appearances,
current: Option<AppearanceIndex>,
}
impl<'a> Iterator for AppearancesIter<'a> {
type Item = AppearanceIndex;
fn next(&mut self) -> Option<AppearanceIndex> {
if let Some(c) = self.current {
self.current = self.appearances[c].next;
Some(c)
} else {
None
}
}
}
impl LocalUseMap {
pub(crate) fn build(
live_locals: &[Local],
elements: &DenseLocationMap,
body: &Body<'_>,
) -> Self {
let nones = IndexVec::from_elem(None, &body.local_decls);
let mut local_use_map = LocalUseMap {
first_def_at: nones.clone(),
first_use_at: nones.clone(),
first_drop_at: nones,
appearances: IndexVec::new(),
};
if live_locals.is_empty() {
return local_use_map;
}
let mut locals_with_use_data: IndexVec<Local, bool> =
IndexVec::from_elem(false, &body.local_decls);
live_locals.iter().for_each(|&local| locals_with_use_data[local] = true);
LocalUseMapBuild { local_use_map: &mut local_use_map, elements, locals_with_use_data }
.visit_body(body);
local_use_map
}
pub(crate) fn defs(&self, local: Local) -> impl Iterator<Item = PointIndex> + '_ {
appearances_iter(self.first_def_at[local], &self.appearances)
.map(move |aa| self.appearances[aa].point_index)
}
pub(crate) fn uses(&self, local: Local) -> impl Iterator<Item = PointIndex> + '_ {
appearances_iter(self.first_use_at[local], &self.appearances)
.map(move |aa| self.appearances[aa].point_index)
}
pub(crate) fn drops(&self, local: Local) -> impl Iterator<Item = PointIndex> + '_ {
appearances_iter(self.first_drop_at[local], &self.appearances)
.map(move |aa| self.appearances[aa].point_index)
}
}
struct LocalUseMapBuild<'me> {
local_use_map: &'me mut LocalUseMap,
elements: &'me DenseLocationMap,
locals_with_use_data: IndexVec<Local, bool>,
}
impl LocalUseMapBuild<'_> {
fn insert_def(&mut self, local: Local, location: Location) {
Self::insert(
self.elements,
&mut self.local_use_map.first_def_at[local],
&mut self.local_use_map.appearances,
location,
);
}
fn insert_use(&mut self, local: Local, location: Location) {
Self::insert(
self.elements,
&mut self.local_use_map.first_use_at[local],
&mut self.local_use_map.appearances,
location,
);
}
fn insert_drop(&mut self, local: Local, location: Location) {
Self::insert(
self.elements,
&mut self.local_use_map.first_drop_at[local],
&mut self.local_use_map.appearances,
location,
);
}
fn insert(
elements: &DenseLocationMap,
first_appearance: &mut Option<AppearanceIndex>,
appearances: &mut Appearances,
location: Location,
) {
let point_index = elements.point_from_location(location);
let appearance_index =
appearances.push(Appearance { point_index, next: *first_appearance });
*first_appearance = Some(appearance_index);
}
}
impl Visitor<'_> for LocalUseMapBuild<'_> {
fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
if self.locals_with_use_data[local] {
match def_use::categorize(context) {
Some(DefUse::Def) => self.insert_def(local, location),
Some(DefUse::Use) => self.insert_use(local, location),
Some(DefUse::Drop) => self.insert_drop(local, location),
_ => (),
}
}
}
}