Auto merge of #141667 - lqd:lazy-maybe-init, r=matthewjasper

Add fast path for maybe-initializedness in liveness

r? `@matthewjasper`

Correct me if I'm wrong Matthew, but my understanding is that
1. `MaybeInitializedPlaces` is currently eagerly computed, in `do_mir_borrowck`
2. but this data is only used in liveness
3. and `liveness::trace` actually only uses it for drop-liveness

This PR moves the computation to `liveness::trace` which looks to be its only use-site. We also add a fast path there, so that it's only computed by drop-liveness.

This is interesting because 1) liveness is only computed for relevant live locals, 2) drop-liveness is only computed for relevant live locals with >0 drop points; 0 is the common case from our benchmarks, as far as I can tell, so even just computing the entire data lazily helps.

It seems possible to also reduce the domain here, and speed up the analysis for the cases where it has to be computed -- so I've left a fixme for that, and may look into it soon.

(I've come upon this while doing implementation work for polonius, so don't be too enamored with possible wins: the goal is to reduce the eventual polonius overhead and make it more palatable 😓)

Author: bors

compiler/rustc_borrowck/src/lib.rs

@@ -40,9 +40,7 @@ use rustc_middle::ty::{
     self, ParamEnv, RegionVid, Ty, TyCtxt, TypeFoldable, TypeVisitable, TypingMode, fold_regions,
 };
 use rustc_middle::{bug, span_bug};
-use rustc_mir_dataflow::impls::{
-    EverInitializedPlaces, MaybeInitializedPlaces, MaybeUninitializedPlaces,
-};
+use rustc_mir_dataflow::impls::{EverInitializedPlaces, MaybeUninitializedPlaces};
 use rustc_mir_dataflow::move_paths::{
     InitIndex, InitLocation, LookupResult, MoveData, MovePathIndex,
 };
@@ -324,10 +322,6 @@ fn do_mir_borrowck<'tcx>(
 
     let move_data = MoveData::gather_moves(body, tcx, |_| true);
 
-    let flow_inits = MaybeInitializedPlaces::new(tcx, body, &move_data)
-        .iterate_to_fixpoint(tcx, body, Some("borrowck"))
-        .into_results_cursor(body);
-
     let locals_are_invalidated_at_exit = tcx.hir_body_owner_kind(def).is_fn_or_closure();
     let borrow_set = BorrowSet::build(tcx, body, locals_are_invalidated_at_exit, &move_data);
 
@@ -346,7 +340,6 @@ fn do_mir_borrowck<'tcx>(
         body,
         &promoted,
         &location_table,
-        flow_inits,
         &move_data,
         &borrow_set,
         consumer_options,

compiler/rustc_borrowck/src/nll.rs

@@ -11,8 +11,6 @@ use rustc_middle::mir::pretty::{PrettyPrintMirOptions, dump_mir_with_options};
 use rustc_middle::mir::{Body, PassWhere, Promoted, create_dump_file, dump_enabled, dump_mir};
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_mir_dataflow::ResultsCursor;
-use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::MoveData;
 use rustc_mir_dataflow::points::DenseLocationMap;
 use rustc_session::config::MirIncludeSpans;
@@ -75,14 +73,13 @@ pub(crate) fn replace_regions_in_mir<'tcx>(
 /// Computes the (non-lexical) regions from the input MIR.
 ///
 /// This may result in errors being reported.
-pub(crate) fn compute_regions<'a, 'tcx>(
+pub(crate) fn compute_regions<'tcx>(
     root_cx: &mut BorrowCheckRootCtxt<'tcx>,
     infcx: &BorrowckInferCtxt<'tcx>,
     universal_regions: UniversalRegions<'tcx>,
     body: &Body<'tcx>,
     promoted: &IndexSlice<Promoted, Body<'tcx>>,
     location_table: &PoloniusLocationTable,
-    flow_inits: ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
     move_data: &MoveData<'tcx>,
     borrow_set: &BorrowSet<'tcx>,
     consumer_options: Option<ConsumerOptions>,
@@ -112,7 +109,6 @@ pub(crate) fn compute_regions<'a, 'tcx>(
         location_table,
         borrow_set,
         &mut polonius_facts,
-        flow_inits,
         move_data,
         Rc::clone(&location_map),
     );

compiler/rustc_borrowck/src/type_check/liveness/mod.rs

@@ -5,8 +5,6 @@ use rustc_middle::mir::{Body, Local, Location, SourceInfo};
 use rustc_middle::span_bug;
 use rustc_middle::ty::relate::Relate;
 use rustc_middle::ty::{GenericArgsRef, Region, RegionVid, Ty, TyCtxt, TypeVisitable};
-use rustc_mir_dataflow::ResultsCursor;
-use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::MoveData;
 use rustc_mir_dataflow::points::DenseLocationMap;
 use tracing::debug;
@@ -28,10 +26,9 @@ mod trace;
 ///
 /// N.B., this computation requires normalization; therefore, it must be
 /// performed before
-pub(super) fn generate<'a, 'tcx>(
+pub(super) fn generate<'tcx>(
     typeck: &mut TypeChecker<'_, 'tcx>,
     location_map: &DenseLocationMap,
-    flow_inits: ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
     move_data: &MoveData<'tcx>,
 ) {
     debug!("liveness::generate");
@@ -58,7 +55,7 @@ pub(super) fn generate<'a, 'tcx>(
     let (relevant_live_locals, boring_locals) =
         compute_relevant_live_locals(typeck.tcx(), &free_regions, typeck.body);
 
-    trace::trace(typeck, location_map, flow_inits, move_data, relevant_live_locals, boring_locals);
+    trace::trace(typeck, location_map, move_data, relevant_live_locals, boring_locals);
 
     // Mark regions that should be live where they appear within rvalues or within a call: like
     // args, regions, and types.

compiler/rustc_borrowck/src/type_check/liveness/trace.rs

@@ -7,10 +7,10 @@ use rustc_middle::mir::{BasicBlock, Body, ConstraintCategory, HasLocalDecls, Loc
 use rustc_middle::traits::query::DropckOutlivesResult;
 use rustc_middle::ty::relate::Relate;
 use rustc_middle::ty::{Ty, TyCtxt, TypeVisitable, TypeVisitableExt};
-use rustc_mir_dataflow::ResultsCursor;
 use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::{HasMoveData, MoveData, MovePathIndex};
 use rustc_mir_dataflow::points::{DenseLocationMap, PointIndex};
+use rustc_mir_dataflow::{Analysis, ResultsCursor};
 use rustc_span::{DUMMY_SP, ErrorGuaranteed, Span};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
 use rustc_trait_selection::traits::ObligationCtxt;
@@ -37,18 +37,17 @@ use crate::type_check::{NormalizeLocation, TypeChecker};
 /// DROP-LIVE set are to the liveness sets for regions found in the
 /// `dropck_outlives` result of the variable's type (in particular,
 /// this respects `#[may_dangle]` annotations).
-pub(super) fn trace<'a, 'tcx>(
+pub(super) fn trace<'tcx>(
     typeck: &mut TypeChecker<'_, 'tcx>,
     location_map: &DenseLocationMap,
-    flow_inits: ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
     move_data: &MoveData<'tcx>,
     relevant_live_locals: Vec<Local>,
     boring_locals: Vec<Local>,
 ) {
     let local_use_map = &LocalUseMap::build(&relevant_live_locals, location_map, typeck.body);
     let cx = LivenessContext {
         typeck,
-        flow_inits,
+        flow_inits: None,
         location_map,
         local_use_map,
         move_data,
@@ -65,7 +64,7 @@ pub(super) fn trace<'a, 'tcx>(
 }
 
 /// Contextual state for the type-liveness coroutine.
-struct LivenessContext<'a, 'typeck, 'b, 'tcx> {
+struct LivenessContext<'a, 'typeck, 'tcx> {
     /// Current type-checker, giving us our inference context etc.
     ///
     /// This also stores the body we're currently analyzing.
@@ -81,8 +80,8 @@ struct LivenessContext<'a, 'typeck, 'b, 'tcx> {
     drop_data: FxIndexMap<Ty<'tcx>, DropData<'tcx>>,
 
     /// Results of dataflow tracking which variables (and paths) have been
-    /// initialized.
-    flow_inits: ResultsCursor<'b, 'tcx, MaybeInitializedPlaces<'b, 'tcx>>,
+    /// initialized. Computed lazily when needed by drop-liveness.
+    flow_inits: Option<ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>>,
 
     /// Index indicating where each variable is assigned, used, or
     /// dropped.
@@ -94,8 +93,8 @@ struct DropData<'tcx> {
     region_constraint_data: Option<&'tcx QueryRegionConstraints<'tcx>>,
 }
 
-struct LivenessResults<'a, 'typeck, 'b, 'tcx> {
-    cx: LivenessContext<'a, 'typeck, 'b, 'tcx>,
+struct LivenessResults<'a, 'typeck, 'tcx> {
+    cx: LivenessContext<'a, 'typeck, 'tcx>,
 
     /// Set of points that define the current local.
     defs: DenseBitSet<PointIndex>,
@@ -116,8 +115,8 @@ struct LivenessResults<'a, 'typeck, 'b, 'tcx> {
     stack: Vec<PointIndex>,
 }
 
-impl<'a, 'typeck, 'b, 'tcx> LivenessResults<'a, 'typeck, 'b, 'tcx> {
-    fn new(cx: LivenessContext<'a, 'typeck, 'b, 'tcx>) -> Self {
+impl<'a, 'typeck, 'tcx> LivenessResults<'a, 'typeck, 'tcx> {
+    fn new(cx: LivenessContext<'a, 'typeck, 'tcx>) -> Self {
         let num_points = cx.location_map.num_points();
         LivenessResults {
             cx,
@@ -459,20 +458,56 @@ impl<'a, 'typeck, 'b, 'tcx> LivenessResults<'a, 'typeck, 'b, 'tcx> {
     }
 }
 
-impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
+impl<'a, 'typeck, 'tcx> LivenessContext<'a, 'typeck, 'tcx> {
+    /// Computes the `MaybeInitializedPlaces` dataflow analysis if it hasn't been done already.
+    ///
+    /// In practice, the results of this dataflow analysis are rarely needed but can be expensive to
+    /// compute on big functions, so we compute them lazily as a fast path when:
+    /// - there are relevant live locals
+    /// - there are drop points for these relevant live locals.
+    ///
+    /// This happens as part of the drop-liveness computation: it's the only place checking for
+    /// maybe-initializedness of `MovePathIndex`es.
+    fn flow_inits(&mut self) -> &mut ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>> {
+        self.flow_inits.get_or_insert_with(|| {
+            let tcx = self.typeck.tcx();
+            let body = self.typeck.body;
+            // FIXME: reduce the `MaybeInitializedPlaces` domain to the useful `MovePath`s.
+            //
+            // This dataflow analysis computes maybe-initializedness of all move paths, which
+            // explains why it can be expensive on big functions. But this data is only used in
+            // drop-liveness. Therefore, most of the move paths computed here are ultimately unused,
+            // even if the results are computed lazily and "no relevant live locals with drop
+            // points" is the common case.
+            //
+            // So we only need the ones for 1) relevant live locals 2) that have drop points. That's
+            // a much, much smaller domain: in our benchmarks, when it's not zero (the most likely
+            // case), there are a few dozens compared to e.g. thousands or tens of thousands of
+            // locals and move paths.
+            let flow_inits = MaybeInitializedPlaces::new(tcx, body, self.move_data)
+                .iterate_to_fixpoint(tcx, body, Some("borrowck"))
+                .into_results_cursor(body);
+            flow_inits
+        })
+    }
+}
+
+impl<'tcx> LivenessContext<'_, '_, 'tcx> {
     fn body(&self) -> &Body<'tcx> {
         self.typeck.body
     }
+
     /// Returns `true` if the local variable (or some part of it) is initialized at the current
     /// cursor position. Callers should call one of the `seek` methods immediately before to point
     /// the cursor to the desired location.
-    fn initialized_at_curr_loc(&self, mpi: MovePathIndex) -> bool {
-        let state = self.flow_inits.get();
+    fn initialized_at_curr_loc(&mut self, mpi: MovePathIndex) -> bool {
+        let flow_inits = self.flow_inits();
+        let state = flow_inits.get();
         if state.contains(mpi) {
             return true;
         }
 
-        let move_paths = &self.flow_inits.analysis().move_data().move_paths;
+        let move_paths = &flow_inits.analysis().move_data().move_paths;
         move_paths[mpi].find_descendant(move_paths, |mpi| state.contains(mpi)).is_some()
     }
 
@@ -481,7 +516,8 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
     /// DROP of some local variable will have an effect -- note that
     /// drops, as they may unwind, are always terminators.
     fn initialized_at_terminator(&mut self, block: BasicBlock, mpi: MovePathIndex) -> bool {
-        self.flow_inits.seek_before_primary_effect(self.body().terminator_loc(block));
+        let terminator_location = self.body().terminator_loc(block);
+        self.flow_inits().seek_before_primary_effect(terminator_location);
         self.initialized_at_curr_loc(mpi)
     }
 
@@ -491,7 +527,8 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
     /// **Warning:** Does not account for the result of `Call`
     /// instructions.
     fn initialized_at_exit(&mut self, block: BasicBlock, mpi: MovePathIndex) -> bool {
-        self.flow_inits.seek_after_primary_effect(self.body().terminator_loc(block));
+        let terminator_location = self.body().terminator_loc(block);
+        self.flow_inits().seek_after_primary_effect(terminator_location);
         self.initialized_at_curr_loc(mpi)
     }
 

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -30,8 +30,6 @@ use rustc_middle::ty::{
     TypeVisitableExt, UserArgs, UserTypeAnnotationIndex, fold_regions,
 };
 use rustc_middle::{bug, span_bug};
-use rustc_mir_dataflow::ResultsCursor;
-use rustc_mir_dataflow::impls::MaybeInitializedPlaces;
 use rustc_mir_dataflow::move_paths::MoveData;
 use rustc_mir_dataflow::points::DenseLocationMap;
 use rustc_span::def_id::CRATE_DEF_ID;
@@ -97,10 +95,9 @@ mod relate_tys;
 /// - `location_table` -- for datalog polonius, the map between `Location`s and `RichLocation`s
 /// - `borrow_set` -- information about borrows occurring in `body`
 /// - `polonius_facts` -- when using Polonius, this is the generated set of Polonius facts
-/// - `flow_inits` -- results of a maybe-init dataflow analysis
 /// - `move_data` -- move-data constructed when performing the maybe-init dataflow analysis
 /// - `location_map` -- map between MIR `Location` and `PointIndex`
-pub(crate) fn type_check<'a, 'tcx>(
+pub(crate) fn type_check<'tcx>(
     root_cx: &mut BorrowCheckRootCtxt<'tcx>,
     infcx: &BorrowckInferCtxt<'tcx>,
     body: &Body<'tcx>,
@@ -109,7 +106,6 @@ pub(crate) fn type_check<'a, 'tcx>(
     location_table: &PoloniusLocationTable,
     borrow_set: &BorrowSet<'tcx>,
     polonius_facts: &mut Option<PoloniusFacts>,
-    flow_inits: ResultsCursor<'a, 'tcx, MaybeInitializedPlaces<'a, 'tcx>>,
     move_data: &MoveData<'tcx>,
     location_map: Rc<DenseLocationMap>,
 ) -> MirTypeckResults<'tcx> {
@@ -167,7 +163,7 @@ pub(crate) fn type_check<'a, 'tcx>(
     typeck.equate_inputs_and_outputs(&normalized_inputs_and_output);
     typeck.check_signature_annotation();
 
-    liveness::generate(&mut typeck, &location_map, flow_inits, move_data);
+    liveness::generate(&mut typeck, &location_map, move_data);
 
     let opaque_type_values =
         opaque_types::take_opaques_and_register_member_constraints(&mut typeck);

tests/mir-opt/dataflow.main.maybe_init.borrowck.dot

@@ -0,0 +1,6 @@
+digraph graph_for_def_id_0_3 {
+    graph[fontname="Courier, monospace"];
+    node[fontname="Courier, monospace"];
+    edge[fontname="Courier, monospace"];
+    bb_0[label=<<table border="1" cellborder="1" cellspacing="0" cellpadding="3" sides="rb"><tr><td colspan="3" sides="tl">bb0</td></tr><tr><td colspan="2" bgcolor="#a0a0a0" sides="tl">MIR</td><td bgcolor="#a0a0a0" sides="tl">STATE</td></tr><tr><td valign="bottom" sides="tl"  align="right"></td><td valign="bottom" sides="tl"  align="left">(on start)</td><td colspan="1" valign="bottom" sides="tl"  align="left">{_0}</td></tr><tr><td valign="top" sides="tl" bgcolor="#f0f0f0" align="right">0</td><td valign="top" sides="tl" bgcolor="#f0f0f0" align="left">_0 = const ()</td><td valign="top" sides="tl" bgcolor="#f0f0f0" align="left"><font color="red">-_0</font></td></tr><tr><td valign="top" sides="tl"  align="right">T</td><td valign="top" sides="tl"  align="left">return</td><td valign="top" sides="tl"  align="left"></td></tr><tr><td valign="bottom" sides="tl" bgcolor="#f0f0f0" align="right"></td><td valign="bottom" sides="tl" bgcolor="#f0f0f0" align="left">(on end)</td><td colspan="1" valign="bottom" sides="tl" bgcolor="#f0f0f0" align="left">{}</td></tr></table>>][shape="none"];
+}

tests/mir-opt/dataflow.main.maybe_uninit.borrowck.dot

@@ -2,5 +2,5 @@
 // Test graphviz dataflow output
 //@ compile-flags: -Z dump-mir=main -Z dump-mir-dataflow
 
-// EMIT_MIR dataflow.main.maybe_init.borrowck.dot
+// EMIT_MIR dataflow.main.maybe_uninit.borrowck.dot
 fn main() {}