Create place and value indices on-demand.

This commit is contained in:
Camille GILLOT 2025-06-22 13:51:39 +00:00 committed by Camille Gillot
parent 051dd12f9e
commit 9415102fb4
5 changed files with 315 additions and 224 deletions

View file

@ -1,3 +1,4 @@
use std::assert_matches::debug_assert_matches;
use std::fmt::{Debug, Formatter};
use std::ops::Range;
@ -350,32 +351,47 @@ pub struct Map<'tcx> {
projections: FxHashMap<(PlaceIndex, TrackElem), PlaceIndex>,
places: IndexVec<PlaceIndex, PlaceInfo<'tcx>>,
value_count: usize,
mode: PlaceCollectionMode,
// The Range corresponds to a slice into `inner_values_buffer`.
inner_values: IndexVec<PlaceIndex, Range<usize>>,
inner_values_buffer: Vec<ValueIndex>,
}
#[derive(Copy, Clone, Debug)]
pub enum PlaceCollectionMode {
Full { value_limit: Option<usize> },
OnDemand,
}
impl<'tcx> Map<'tcx> {
/// Returns a map that only tracks places whose type has scalar layout.
///
/// This is currently the only way to create a [`Map`]. The way in which the tracked places are
/// chosen is an implementation detail and may not be relied upon (other than that their type
/// are scalars).
pub fn new(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, value_limit: Option<usize>) -> Self {
let capacity = 4 * body.local_decls.len() + value_limit.unwrap_or(body.local_decls.len());
#[tracing::instrument(level = "trace", skip(tcx, body))]
pub fn new(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, mode: PlaceCollectionMode) -> Self {
tracing::trace!(def_id=?body.source.def_id());
let capacity = 4 * body.local_decls.len();
let mut map = Self {
locals: IndexVec::from_elem(None, &body.local_decls),
projections: FxHashMap::default(),
places: IndexVec::with_capacity(capacity),
value_count: 0,
inner_values: IndexVec::with_capacity(capacity),
mode,
inner_values: IndexVec::new(),
inner_values_buffer: Vec::new(),
};
map.register_locals(tcx, body);
map.collect_places(tcx, body);
map.propagate_assignments(tcx, body);
map.create_values(tcx, body, value_limit);
map.trim_useless_places();
match mode {
PlaceCollectionMode::Full { value_limit } => {
map.collect_places(tcx, body);
map.propagate_assignments(tcx, body);
map.create_values(tcx, body, value_limit);
map.trim_useless_places();
}
PlaceCollectionMode::OnDemand => {}
}
debug!("registered {} places ({} nodes in total)", map.value_count, map.places.len());
map
}
@ -427,12 +443,18 @@ impl<'tcx> Map<'tcx> {
match rhs {
Rvalue::Use(Operand::Move(rhs) | Operand::Copy(rhs))
| Rvalue::CopyForDeref(rhs) => {
let Some(lhs) = self.register_place(tcx, body, *lhs) else { continue };
let Some(rhs) = self.register_place(tcx, body, *rhs) else { continue };
let Some(lhs) = self.register_place_and_discr(tcx, body, *lhs) else {
continue;
};
let Some(rhs) = self.register_place_and_discr(tcx, body, *rhs) else {
continue;
};
assignments.insert((lhs, rhs));
}
Rvalue::Aggregate(kind, fields) => {
let Some(mut lhs) = self.register_place(tcx, body, *lhs) else { continue };
let Some(mut lhs) = self.register_place_and_discr(tcx, body, *lhs) else {
continue;
};
match **kind {
// Do not propagate unions.
AggregateKind::Adt(_, _, _, _, Some(_)) => continue,
@ -455,7 +477,7 @@ impl<'tcx> Map<'tcx> {
}
for (index, field) in fields.iter_enumerated() {
if let Some(rhs) = field.place()
&& let Some(rhs) = self.register_place(tcx, body, rhs)
&& let Some(rhs) = self.register_place_and_discr(tcx, body, rhs)
{
let lhs = self.register_place_index(
self.places[rhs].ty,
@ -512,6 +534,7 @@ impl<'tcx> Map<'tcx> {
/// Create values for places whose type have scalar layout.
#[tracing::instrument(level = "trace", skip(self, tcx, body))]
fn create_values(&mut self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>, value_limit: Option<usize>) {
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
let typing_env = body.typing_env(tcx);
for place_info in self.places.iter_mut() {
// The user requires a bound on the number of created values.
@ -550,6 +573,7 @@ impl<'tcx> Map<'tcx> {
/// Trim useless places.
#[tracing::instrument(level = "trace", skip(self))]
fn trim_useless_places(&mut self) {
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
for opt_place in self.locals.iter_mut() {
if let Some(place) = *opt_place
&& self.inner_values[place].is_empty()
@ -562,7 +586,7 @@ impl<'tcx> Map<'tcx> {
}
#[tracing::instrument(level = "trace", skip(self), ret)]
fn register_place_index(
pub fn register_place_index(
&mut self,
ty: Ty<'tcx>,
base: PlaceIndex,
@ -576,49 +600,124 @@ impl<'tcx> Map<'tcx> {
})
}
#[tracing::instrument(level = "trace", skip(self, tcx, body))]
fn register_place(
#[tracing::instrument(level = "trace", skip(self, tcx, body), ret)]
pub fn register_place(
&mut self,
tcx: TyCtxt<'tcx>,
body: &Body<'tcx>,
place: Place<'tcx>,
tail: Option<TrackElem>,
) -> Option<PlaceIndex> {
// Create a place for this projection.
let mut place_index = self.locals[place.local]?;
let mut ty = PlaceTy::from_ty(body.local_decls[place.local].ty);
tracing::trace!(?place_index, ?ty);
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.ty.kind()
&& let ty::Slice(..) = ref_ty.kind()
{
self.register_place_index(tcx.types.usize, place_index, TrackElem::DerefLen);
} else if ty.ty.is_enum() {
let discriminant_ty = ty.ty.discriminant_ty(tcx);
self.register_place_index(discriminant_ty, place_index, TrackElem::Discriminant);
}
for proj in place.projection {
let track_elem = proj.try_into().ok()?;
ty = ty.projection_ty(tcx, proj);
place_index = self.register_place_index(ty.ty, place_index, track_elem);
tracing::trace!(?proj, ?place_index, ?ty);
}
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.ty.kind()
&& let ty::Slice(..) = ref_ty.kind()
{
self.register_place_index(tcx.types.usize, place_index, TrackElem::DerefLen);
} else if ty.ty.is_enum() {
let discriminant_ty = ty.ty.discriminant_ty(tcx);
self.register_place_index(discriminant_ty, place_index, TrackElem::Discriminant);
}
if let Some(tail) = tail {
let ty = match tail {
TrackElem::Discriminant => ty.ty.discriminant_ty(tcx),
TrackElem::Variant(..) | TrackElem::Field(..) => todo!(),
TrackElem::DerefLen => tcx.types.usize,
};
place_index = self.register_place_index(ty, place_index, tail);
}
Some(place_index)
}
#[tracing::instrument(level = "trace", skip(self, tcx, body), ret)]
fn register_place_and_discr(
&mut self,
tcx: TyCtxt<'tcx>,
body: &Body<'tcx>,
place: Place<'tcx>,
) -> Option<PlaceIndex> {
let place = self.register_place(tcx, body, place, None)?;
let ty = self.places[place].ty;
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.kind()
&& let ty::Slice(..) = ref_ty.kind()
{
self.register_place_index(tcx.types.usize, place, TrackElem::DerefLen);
} else if ty.is_enum() {
let discriminant_ty = ty.discriminant_ty(tcx);
self.register_place_index(discriminant_ty, place, TrackElem::Discriminant);
}
Some(place)
}
#[tracing::instrument(level = "trace", skip(self, tcx, typing_env), ret)]
pub fn register_value(
&mut self,
tcx: TyCtxt<'tcx>,
typing_env: ty::TypingEnv<'tcx>,
place: PlaceIndex,
) -> Option<ValueIndex> {
let place_info = &mut self.places[place];
if let Some(value) = place_info.value_index {
return Some(value);
}
if let Ok(ty) = tcx.try_normalize_erasing_regions(typing_env, place_info.ty) {
place_info.ty = ty;
}
// Allocate a value slot if it doesn't have one, and the user requested one.
if let Ok(layout) = tcx.layout_of(typing_env.as_query_input(place_info.ty))
&& layout.backend_repr.is_scalar()
{
place_info.value_index = Some(self.value_count.into());
self.value_count += 1;
}
place_info.value_index
}
#[tracing::instrument(level = "trace", skip(self, f))]
pub fn register_copy_tree(
&mut self,
// Tree to copy.
source: PlaceIndex,
// Tree to build.
target: PlaceIndex,
f: &mut impl FnMut(ValueIndex, ValueIndex),
) {
if let Some(source_value) = self.places[source].value_index {
let target_value = *self.places[target].value_index.get_or_insert_with(|| {
let value_index = self.value_count.into();
self.value_count += 1;
value_index
});
f(source_value, target_value)
}
// Iterate over `source` children and recurse.
let mut source_child_iter = self.places[source].first_child;
while let Some(source_child) = source_child_iter {
source_child_iter = self.places[source_child].next_sibling;
// Try to find corresponding child and recurse. Reasoning is similar as above.
let source_info = &self.places[source_child];
let source_ty = source_info.ty;
let source_elem = source_info.proj_elem.unwrap();
let target_child = self.register_place_index(source_ty, target, source_elem);
self.register_copy_tree(source_child, target_child, f);
}
}
/// Precompute the list of values inside `root` and store it inside
/// as a slice within `inner_values_buffer`.
#[tracing::instrument(level = "trace", skip(self))]
fn cache_preorder_invoke(&mut self, root: PlaceIndex) {
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
let start = self.inner_values_buffer.len();
if let Some(vi) = self.places[root].value_index {
self.inner_values_buffer.push(vi);
@ -649,7 +748,7 @@ impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, 'tcx> {
return;
}
self.map.register_place(self.tcx, self.body, *place);
self.map.register_place_and_discr(self.tcx, self.body, *place);
}
}
@ -724,6 +823,7 @@ impl<'tcx> Map<'tcx> {
///
/// `tail_elem` allows to support discriminants that are not a place in MIR, but that we track
/// as such.
#[tracing::instrument(level = "trace", skip(self, f))]
pub fn for_each_aliasing_place(
&self,
place: PlaceRef<'_>,
@ -758,6 +858,7 @@ impl<'tcx> Map<'tcx> {
}
/// Invoke the given function on all the descendants of the given place, except one branch.
#[tracing::instrument(level = "trace", skip(self, f))]
fn for_each_variant_sibling(
&self,
parent: PlaceIndex,
@ -784,9 +885,22 @@ impl<'tcx> Map<'tcx> {
}
/// Invoke a function on each value in the given place and all descendants.
#[tracing::instrument(level = "trace", skip(self, f))]
fn for_each_value_inside(&self, root: PlaceIndex, f: &mut impl FnMut(ValueIndex)) {
for &v in self.values_inside(root) {
f(v)
if let Some(range) = self.inner_values.get(root) {
// Optimized path: we have cached the inner values.
let values = &self.inner_values_buffer[range.clone()];
for &v in values {
f(v)
}
} else {
if let Some(root) = self.places[root].value_index {
f(root)
}
for child in self.children(root) {
self.for_each_value_inside(child, f);
}
}
}
@ -799,7 +913,9 @@ impl<'tcx> Map<'tcx> {
f: &mut impl FnMut(PlaceIndex, &O),
) {
// Fast path is there is nothing to do.
if self.inner_values[root].is_empty() {
if let Some(value_range) = self.inner_values.get(root)
&& value_range.is_empty()
{
return;
}