Skip to content

Commit 42dc029

Browse files
committed
Create place and value indices on-demand.
1 parent 434a99c commit 42dc029

File tree

3 files changed

+210
-72
lines changed

3 files changed

+210
-72
lines changed

compiler/rustc_mir_dataflow/src/value_analysis.rs

Lines changed: 151 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
use std::assert_matches::debug_assert_matches;
12
use std::fmt::{Debug, Formatter};
23
use std::ops::Range;
34

@@ -350,32 +351,47 @@ pub struct Map<'tcx> {
350351
projections: FxHashMap<(PlaceIndex, TrackElem), PlaceIndex>,
351352
places: IndexVec<PlaceIndex, PlaceInfo<'tcx>>,
352353
value_count: usize,
354+
mode: PlaceCollectionMode,
353355
// The Range corresponds to a slice into `inner_values_buffer`.
354356
inner_values: IndexVec<PlaceIndex, Range<usize>>,
355357
inner_values_buffer: Vec<ValueIndex>,
356358
}
357359

360+
#[derive(Copy, Clone, Debug)]
361+
pub enum PlaceCollectionMode {
362+
Full { value_limit: Option<usize> },
363+
OnDemand,
364+
}
365+
358366
impl<'tcx> Map<'tcx> {
359367
/// Returns a map that only tracks places whose type has scalar layout.
360368
///
361369
/// This is currently the only way to create a [`Map`]. The way in which the tracked places are
362370
/// chosen is an implementation detail and may not be relied upon (other than that their type
363371
/// are scalars).
364-
pub fn new(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, value_limit: Option<usize>) -> Self {
365-
let capacity = 4 * body.local_decls.len() + value_limit.unwrap_or(body.local_decls.len());
372+
#[tracing::instrument(level = "trace", skip(tcx, body))]
373+
pub fn new(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, mode: PlaceCollectionMode) -> Self {
374+
tracing::trace!(def_id=?body.source.def_id());
375+
let capacity = 4 * body.local_decls.len();
366376
let mut map = Self {
367377
locals: IndexVec::from_elem(None, &body.local_decls),
368378
projections: FxHashMap::default(),
369379
places: IndexVec::with_capacity(capacity),
370380
value_count: 0,
371-
inner_values: IndexVec::with_capacity(capacity),
381+
mode,
382+
inner_values: IndexVec::new(),
372383
inner_values_buffer: Vec::new(),
373384
};
374385
map.register_locals(tcx, body);
375-
map.collect_places(tcx, body);
376-
map.propagate_assignments(tcx, body);
377-
map.create_values(tcx, body, value_limit);
378-
map.trim_useless_places();
386+
match mode {
387+
PlaceCollectionMode::Full { value_limit } => {
388+
map.collect_places(tcx, body);
389+
map.propagate_assignments(tcx, body);
390+
map.create_values(tcx, body, value_limit);
391+
map.trim_useless_places();
392+
}
393+
PlaceCollectionMode::OnDemand => {}
394+
}
379395
debug!("registered {} places ({} nodes in total)", map.value_count, map.places.len());
380396
map
381397
}
@@ -427,12 +443,18 @@ impl<'tcx> Map<'tcx> {
427443
match rhs {
428444
Rvalue::Use(Operand::Move(rhs) | Operand::Copy(rhs))
429445
| Rvalue::CopyForDeref(rhs) => {
430-
let Some(lhs) = self.register_place(tcx, body, *lhs) else { continue };
431-
let Some(rhs) = self.register_place(tcx, body, *rhs) else { continue };
446+
let Some(lhs) = self.register_place_and_discr(tcx, body, *lhs) else {
447+
continue;
448+
};
449+
let Some(rhs) = self.register_place_and_discr(tcx, body, *rhs) else {
450+
continue;
451+
};
432452
assignments.insert((lhs, rhs));
433453
}
434454
Rvalue::Aggregate(kind, fields) => {
435-
let Some(mut lhs) = self.register_place(tcx, body, *lhs) else { continue };
455+
let Some(mut lhs) = self.register_place_and_discr(tcx, body, *lhs) else {
456+
continue;
457+
};
436458
match **kind {
437459
// Do not propagate unions.
438460
AggregateKind::Adt(_, _, _, _, Some(_)) => continue,
@@ -455,7 +477,7 @@ impl<'tcx> Map<'tcx> {
455477
}
456478
for (index, field) in fields.iter_enumerated() {
457479
if let Some(rhs) = field.place()
458-
&& let Some(rhs) = self.register_place(tcx, body, rhs)
480+
&& let Some(rhs) = self.register_place_and_discr(tcx, body, rhs)
459481
{
460482
let lhs = self.register_place_index(
461483
self.places[rhs].ty,
@@ -512,6 +534,7 @@ impl<'tcx> Map<'tcx> {
512534
/// Create values for places whose type have scalar layout.
513535
#[tracing::instrument(level = "trace", skip(self, tcx, body))]
514536
fn create_values(&mut self, tcx: TyCtxt<'tcx>, body: &Body<'tcx>, value_limit: Option<usize>) {
537+
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
515538
let typing_env = body.typing_env(tcx);
516539
for place_info in self.places.iter_mut() {
517540
// The user requires a bound on the number of created values.
@@ -550,6 +573,7 @@ impl<'tcx> Map<'tcx> {
550573
/// Trim useless places.
551574
#[tracing::instrument(level = "trace", skip(self))]
552575
fn trim_useless_places(&mut self) {
576+
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
553577
for opt_place in self.locals.iter_mut() {
554578
if let Some(place) = *opt_place
555579
&& self.inner_values[place].is_empty()
@@ -562,7 +586,7 @@ impl<'tcx> Map<'tcx> {
562586
}
563587

564588
#[tracing::instrument(level = "trace", skip(self), ret)]
565-
fn register_place_index(
589+
pub fn register_place_index(
566590
&mut self,
567591
ty: Ty<'tcx>,
568592
base: PlaceIndex,
@@ -576,49 +600,124 @@ impl<'tcx> Map<'tcx> {
576600
})
577601
}
578602

579-
#[tracing::instrument(level = "trace", skip(self, tcx, body))]
580-
fn register_place(
603+
#[tracing::instrument(level = "trace", skip(self, tcx, body), ret)]
604+
pub fn register_place(
581605
&mut self,
582606
tcx: TyCtxt<'tcx>,
583607
body: &Body<'tcx>,
584608
place: Place<'tcx>,
609+
tail: Option<TrackElem>,
585610
) -> Option<PlaceIndex> {
586611
// Create a place for this projection.
587612
let mut place_index = self.locals[place.local]?;
588613
let mut ty = PlaceTy::from_ty(body.local_decls[place.local].ty);
589614
tracing::trace!(?place_index, ?ty);
590615

591-
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.ty.kind()
592-
&& let ty::Slice(..) = ref_ty.kind()
593-
{
594-
self.register_place_index(tcx.types.usize, place_index, TrackElem::DerefLen);
595-
} else if ty.ty.is_enum() {
596-
let discriminant_ty = ty.ty.discriminant_ty(tcx);
597-
self.register_place_index(discriminant_ty, place_index, TrackElem::Discriminant);
598-
}
599-
600616
for proj in place.projection {
601617
let track_elem = proj.try_into().ok()?;
602618
ty = ty.projection_ty(tcx, proj);
603619
place_index = self.register_place_index(ty.ty, place_index, track_elem);
604620
tracing::trace!(?proj, ?place_index, ?ty);
621+
}
605622

606-
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.ty.kind()
607-
&& let ty::Slice(..) = ref_ty.kind()
608-
{
609-
self.register_place_index(tcx.types.usize, place_index, TrackElem::DerefLen);
610-
} else if ty.ty.is_enum() {
611-
let discriminant_ty = ty.ty.discriminant_ty(tcx);
612-
self.register_place_index(discriminant_ty, place_index, TrackElem::Discriminant);
613-
}
623+
if let Some(tail) = tail {
624+
let ty = match tail {
625+
TrackElem::Discriminant => ty.ty.discriminant_ty(tcx),
626+
TrackElem::Variant(..) | TrackElem::Field(..) => todo!(),
627+
TrackElem::DerefLen => tcx.types.usize,
628+
};
629+
place_index = self.register_place_index(ty, place_index, tail);
614630
}
615631

616632
Some(place_index)
617633
}
618634

635+
#[tracing::instrument(level = "trace", skip(self, tcx, body), ret)]
636+
fn register_place_and_discr(
637+
&mut self,
638+
tcx: TyCtxt<'tcx>,
639+
body: &Body<'tcx>,
640+
place: Place<'tcx>,
641+
) -> Option<PlaceIndex> {
642+
let place = self.register_place(tcx, body, place, None)?;
643+
let ty = self.places[place].ty;
644+
645+
if let ty::Ref(_, ref_ty, _) | ty::RawPtr(ref_ty, _) = ty.kind()
646+
&& let ty::Slice(..) = ref_ty.kind()
647+
{
648+
self.register_place_index(tcx.types.usize, place, TrackElem::DerefLen);
649+
} else if ty.is_enum() {
650+
let discriminant_ty = ty.discriminant_ty(tcx);
651+
self.register_place_index(discriminant_ty, place, TrackElem::Discriminant);
652+
}
653+
654+
Some(place)
655+
}
656+
657+
#[tracing::instrument(level = "trace", skip(self, tcx, typing_env), ret)]
658+
pub fn register_value(
659+
&mut self,
660+
tcx: TyCtxt<'tcx>,
661+
typing_env: ty::TypingEnv<'tcx>,
662+
place: PlaceIndex,
663+
) -> Option<ValueIndex> {
664+
let place_info = &mut self.places[place];
665+
if let Some(value) = place_info.value_index {
666+
return Some(value);
667+
}
668+
669+
if let Ok(ty) = tcx.try_normalize_erasing_regions(typing_env, place_info.ty) {
670+
place_info.ty = ty;
671+
}
672+
673+
// Allocate a value slot if it doesn't have one, and the user requested one.
674+
if let Ok(layout) = tcx.layout_of(typing_env.as_query_input(place_info.ty))
675+
&& layout.backend_repr.is_scalar()
676+
{
677+
place_info.value_index = Some(self.value_count.into());
678+
self.value_count += 1;
679+
}
680+
681+
place_info.value_index
682+
}
683+
684+
#[tracing::instrument(level = "trace", skip(self, f))]
685+
pub fn register_copy_tree(
686+
&mut self,
687+
// Tree to copy.
688+
source: PlaceIndex,
689+
// Tree to build.
690+
target: PlaceIndex,
691+
f: &mut impl FnMut(ValueIndex, ValueIndex),
692+
) {
693+
if let Some(source_value) = self.places[source].value_index {
694+
let target_value = *self.places[target].value_index.get_or_insert_with(|| {
695+
let value_index = self.value_count.into();
696+
self.value_count += 1;
697+
value_index
698+
});
699+
f(source_value, target_value)
700+
}
701+
702+
// Iterate over `source` children and recurse.
703+
let mut source_child_iter = self.places[source].first_child;
704+
while let Some(source_child) = source_child_iter {
705+
source_child_iter = self.places[source_child].next_sibling;
706+
707+
// Try to find corresponding child and recurse. Reasoning is similar as above.
708+
let source_info = &self.places[source_child];
709+
let source_ty = source_info.ty;
710+
let source_elem = source_info.proj_elem.unwrap();
711+
let target_child = self.register_place_index(source_ty, target, source_elem);
712+
self.register_copy_tree(source_child, target_child, f);
713+
}
714+
}
715+
619716
/// Precompute the list of values inside `root` and store it inside
620717
/// as a slice within `inner_values_buffer`.
718+
#[tracing::instrument(level = "trace", skip(self))]
621719
fn cache_preorder_invoke(&mut self, root: PlaceIndex) {
720+
debug_assert_matches!(self.mode, PlaceCollectionMode::Full { .. });
622721
let start = self.inner_values_buffer.len();
623722
if let Some(vi) = self.places[root].value_index {
624723
self.inner_values_buffer.push(vi);
@@ -649,7 +748,7 @@ impl<'tcx> Visitor<'tcx> for PlaceCollector<'_, 'tcx> {
649748
return;
650749
}
651750

652-
self.map.register_place(self.tcx, self.body, *place);
751+
self.map.register_place_and_discr(self.tcx, self.body, *place);
653752
}
654753
}
655754

@@ -724,6 +823,7 @@ impl<'tcx> Map<'tcx> {
724823
///
725824
/// `tail_elem` allows to support discriminants that are not a place in MIR, but that we track
726825
/// as such.
826+
#[tracing::instrument(level = "trace", skip(self, f))]
727827
pub fn for_each_aliasing_place(
728828
&self,
729829
place: PlaceRef<'_>,
@@ -761,6 +861,7 @@ impl<'tcx> Map<'tcx> {
761861
}
762862

763863
/// Invoke the given function on all the descendants of the given place, except one branch.
864+
#[tracing::instrument(level = "trace", skip(self, f))]
764865
fn for_each_variant_sibling(
765866
&self,
766867
parent: PlaceIndex,
@@ -781,11 +882,22 @@ impl<'tcx> Map<'tcx> {
781882
}
782883

783884
/// Invoke a function on each value in the given place and all descendants.
885+
#[tracing::instrument(level = "trace", skip(self, f))]
784886
fn for_each_value_inside(&self, root: PlaceIndex, f: &mut impl FnMut(ValueIndex)) {
785-
let range = self.inner_values[root].clone();
786-
let values = &self.inner_values_buffer[range];
787-
for &v in values {
788-
f(v)
887+
if let Some(range) = self.inner_values.get(root) {
888+
// Optimized path: we have cached the inner values.
889+
let values = &self.inner_values_buffer[range.clone()];
890+
for &v in values {
891+
f(v)
892+
}
893+
} else {
894+
if let Some(root) = self.places[root].value_index {
895+
f(root)
896+
}
897+
898+
for child in self.children(root) {
899+
self.for_each_value_inside(child, f);
900+
}
789901
}
790902
}
791903

@@ -798,7 +910,9 @@ impl<'tcx> Map<'tcx> {
798910
f: &mut impl FnMut(PlaceIndex, &O),
799911
) {
800912
// Fast path is there is nothing to do.
801-
if self.inner_values[root].is_empty() {
913+
if let Some(value_range) = self.inner_values.get(root)
914+
&& value_range.is_empty()
915+
{
802916
return;
803917
}
804918

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use rustc_middle::ty::{self, Ty, TyCtxt};
2020
use rustc_mir_dataflow::fmt::DebugWithContext;
2121
use rustc_mir_dataflow::lattice::{FlatSet, HasBottom};
2222
use rustc_mir_dataflow::value_analysis::{
23-
Map, PlaceIndex, State, TrackElem, ValueOrPlace, debug_with_context,
23+
Map, PlaceCollectionMode, PlaceIndex, State, TrackElem, ValueOrPlace, debug_with_context,
2424
};
2525
use rustc_mir_dataflow::{Analysis, ResultsVisitor, visit_reachable_results};
2626
use rustc_span::DUMMY_SP;
@@ -54,10 +54,10 @@ impl<'tcx> crate::MirPass<'tcx> for DataflowConstProp {
5454
// `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of
5555
// map nodes is strongly correlated to the number of tracked places, this becomes more or
5656
// less `O(n)` if we place a constant limit on the number of tracked places.
57-
let place_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
57+
let value_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
5858

5959
// Decide which places to track during the analysis.
60-
let map = Map::new(tcx, body, place_limit);
60+
let map = Map::new(tcx, body, PlaceCollectionMode::Full { value_limit });
6161

6262
// Perform the actual dataflow analysis.
6363
let mut const_ = debug_span!("analyze")

0 commit comments

Comments
 (0)