Create place and value indices on-demand.

cjgillot · cjgillot · commit c12e371ec8b5 · 2025-07-02T10:56:26.000Z
diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -1,3 +1,4 @@
+use std::assert_matches::debug_assert_matches;
 use std::fmt::{Debug, Formatter};
 use std::ops::Range;
 
@@ -352,32 +353,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
     }
@@ -429,12 +445,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,
@@ -457,7 +479,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,
@@ -514,6 +536,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.
@@ -552,6 +575,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()
@@ -564,7 +588,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,
@@ -578,49 +602,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);
@@ -651,7 +750,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);
     }
 }
 
@@ -726,6 +825,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<'_>,
@@ -763,6 +863,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,
@@ -783,11 +884,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)) {
-        let range = self.inner_values[root].clone();
-        let values = &self.inner_values_buffer[range];
-        for &v in values {
-            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);
+            }
         }
     }
 
@@ -800,7 +912,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;
         }
 
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -21,7 +21,7 @@ use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_mir_dataflow::fmt::DebugWithContext;
 use rustc_mir_dataflow::lattice::{FlatSet, HasBottom};
 use rustc_mir_dataflow::value_analysis::{
-    Map, PlaceIndex, State, TrackElem, ValueOrPlace, debug_with_context,
+    Map, PlaceCollectionMode, PlaceIndex, State, TrackElem, ValueOrPlace, debug_with_context,
 };
 use rustc_mir_dataflow::{Analysis, ResultsVisitor, visit_reachable_results};
 use rustc_span::DUMMY_SP;
@@ -55,10 +55,10 @@ impl<'tcx> crate::MirPass<'tcx> for DataflowConstProp {
         // `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of
         // map nodes is strongly correlated to the number of tracked places, this becomes more or
         // less `O(n)` if we place a constant limit on the number of tracked places.
-        let place_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
+        let value_limit = if tcx.sess.mir_opt_level() < 4 { Some(PLACE_LIMIT) } else { None };
 
         // Decide which places to track during the analysis.
-        let map = Map::new(tcx, body, place_limit);
+        let map = Map::new(tcx, body, PlaceCollectionMode::Full { value_limit });
 
         // Perform the actual dataflow analysis.
         let mut const_ = debug_span!("analyze")
diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
