Speed up EverInitializedPlaces

compiler/rustc_borrowck/src/lib.rs

@@ -2497,13 +2497,14 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
         // partial initialization, do not complain about mutability
         // errors except for actual mutation (as opposed to an attempt
         // to do a partial initialization).
-        let previously_initialized = self.is_local_ever_initialized(place.local, state);
+        let previously_initialized = state.ever_inits.contains(place.local);
 
         // at this point, we have set up the error reporting state.
-        if let Some(init_index) = previously_initialized {
+        if previously_initialized {
             if let (AccessKind::Mutate, Some(_)) = (error_access, place.as_local()) {
                 // If this is a mutate access to an immutable local variable with no projections
                 // report the error as an illegal reassignment
+                let init_index = self.first_reaching_init(place.local, location).unwrap();
                 let init = &self.move_data.inits[init_index];
                 let assigned_span = init.span(self.body);
                 self.report_illegal_reassignment((place, span), assigned_span, place);
@@ -2516,10 +2517,14 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
         }
     }
 
-    fn is_local_ever_initialized(&self, local: Local, state: &BorrowckDomain) -> Option<InitIndex> {
+    /// Returns the first init of `local` (in gather order) that may have executed on some path
+    /// reaching `location` without an intervening `StorageDead(local)`.
+    fn first_reaching_init(&self, local: Local, location: Location) -> Option<InitIndex> {
         let mpi = self.move_data.rev_lookup.find_local(local)?;
-        let ii = &self.move_data.init_path_map[mpi];
-        ii.into_iter().find(|&&index| state.ever_inits.contains(index)).copied()
+        self.move_data.init_path_map[mpi].iter().copied().find(|&ii| {
+            let init = self.move_data.inits[ii];
+            EverInitializedPlaces::init_reaches_location(self.body, local, init, location)
+        })
     }
 
     /// Adds the place into the used mutable variables set
@@ -2530,7 +2535,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
                 // mutated, then it is justified to be annotated with the `mut`
                 // keyword, since the mutation may be a possible reassignment.
                 if is_local_mutation_allowed != LocalMutationIsAllowed::Yes
-                    && self.is_local_ever_initialized(local, state).is_some()
+                    && state.ever_inits.contains(local)
                 {
                     self.used_mut.insert(local);
                 }

compiler/rustc_mir_dataflow/src/impls/initialized.rs

@@ -1,16 +1,20 @@
 use std::assert_matches;
 
 use rustc_abi::VariantIdx;
-use rustc_index::Idx;
+use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_middle::bug;
-use rustc_middle::mir::{self, Body, CallReturnPlaces, Location, TerminatorEdges};
+use rustc_middle::mir::{
+    self, BasicBlock, Body, CallReturnPlaces, Local, Location, StatementKind, TerminatorEdges,
+};
 use rustc_middle::ty::{self, TyCtxt};
 use smallvec::SmallVec;
-use tracing::{debug, instrument};
+use tracing::instrument;
 
 use crate::drop_flag_effects::{DropFlagState, InactiveVariants};
-use crate::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex};
+use crate::move_paths::{
+    HasMoveData, Init, InitKind, InitLocation, LookupResult, MoveData, MovePathIndex,
+};
 use crate::{
     Analysis, GenKill, MaybeReachable, SwitchTargetIndex, drop_flag_effects,
     drop_flag_effects_for_function_entry, drop_flag_effects_for_location, on_all_children_bits,
@@ -267,36 +271,34 @@ impl<'tcx> HasMoveData<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
     }
 }
 
-/// `EverInitializedPlaces` tracks all initializations that may have occurred
-/// upon reaching a particular point in the control flow for a function,
-/// without an intervening `StorageDead`.
+/// `EverInitializedPlaces` tracks all initializations of locals that may have
+/// occurred upon reaching a particular point in the control flow for a
+/// function, without an intervening `StorageDead`.
 ///
 /// This dataflow is used to determine if an immutable local variable may
 /// be assigned to.
 ///
 /// For example, in code like the following, we have corresponding
-/// dataflow information shown in the right-hand comments. Underscored indices
-/// are used to distinguish between multiple initializations of the same local
-/// variable, e.g. `b_0` and `b_1`.
+/// dataflow information shown in the right-hand comments.
 ///
 /// ```rust
 /// struct S;
 /// #[rustfmt::skip]
 /// fn foo(p: bool) {                           // ever-init:
-///                                             // {p,                  }
-///     let a = S; let mut b = S; let c; let d; // {p, a, b_0,          }
+///                                             // {p,           }
+///     let a = S; let mut b = S; let c; let d; // {p, a, b,     }
 ///
 ///     if p {
-///         drop(a);                            // {p, a, b_0,          }
-///         b = S;                              // {p, a, b_0, b_1,     }
+///         drop(a);                            // {p, a, b,     }
+///         b = S;                              // {p, a, b,     }
 ///
 ///     } else {
-///         drop(b);                            // {p, a, b_0, b_1,     }
-///         d = S;                              // {p, a, b_0, b_1,    d}
+///         drop(b);                            // {p, a, b,     }
+///         d = S;                              // {p, a, b,    d}
 ///
-///     }                                       // {p, a, b_0, b_1,    d}
+///     }                                       // {p, a, b,    d}
 ///
-///     c = S;                                  // {p, a, b_0, b_1, c, d}
+///     c = S;                                  // {p, a, b, c, d}
 /// }
 /// ```
 pub struct EverInitializedPlaces<'a, 'tcx> {
@@ -590,23 +592,21 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
     }
 }
 
-/// There can be many more `InitIndex` than there are locals in a MIR body.
-/// We use a mixed bitset to avoid paying too high a memory footprint.
-pub type EverInitializedPlacesDomain = MixedBitSet<InitIndex>;
+pub type EverInitializedPlacesDomain = DenseBitSet<Local>;
 
 impl<'tcx> Analysis<'tcx> for EverInitializedPlaces<'_, 'tcx> {
     type Domain = EverInitializedPlacesDomain;
 
     const NAME: &'static str = "ever_init";
 
-    fn bottom_value(&self, _: &mir::Body<'tcx>) -> Self::Domain {
-        // bottom = no initialized variables by default
-        MixedBitSet::new_empty(self.move_data().inits.len())
+    fn bottom_value(&self, body: &mir::Body<'tcx>) -> Self::Domain {
+        // bottom = no initialized locals by default
+        DenseBitSet::new_empty(body.local_decls.len())
     }
 
     fn initialize_start_block(&self, body: &mir::Body<'tcx>, state: &mut Self::Domain) {
-        for arg_init in 0..body.arg_count {
-            state.insert(InitIndex::new(arg_init));
+        for arg in body.args_iter() {
+            state.insert(arg);
         }
     }
 
@@ -618,20 +618,18 @@ impl<'tcx> Analysis<'tcx> for EverInitializedPlaces<'_, 'tcx> {
         location: Location,
     ) {
         let move_data = self.move_data();
-        let init_path_map = &move_data.init_path_map;
         let init_loc_map = &move_data.init_loc_map;
-        let rev_lookup = &move_data.rev_lookup;
 
-        debug!("initializes move_indexes {:?}", init_loc_map[location]);
-        state.gen_all(init_loc_map[location].iter().copied());
+        // Record inits of locals. Projections can be ignored.
+        state.gen_all(init_loc_map[location].iter().copied().filter_map(|ii| {
+            let init_mpi = move_data.inits[ii].path;
+            move_data.move_paths[init_mpi].place.as_local()
+        }));
 
-        if let mir::StatementKind::StorageDead(local) = stmt.kind
-            // End inits for StorageDead, so that an immutable variable can
-            // be reinitialized on the next iteration of the loop.
-            && let Some(move_path_index) = rev_lookup.find_local(local)
-        {
-            debug!("clears the ever initialized status of {:?}", init_path_map[move_path_index]);
-            state.kill_all(init_path_map[move_path_index].iter().copied());
+        // Kill on StorageDead, so that an immutable variable can
+        // be reinitialized on the next iteration of the loop.
+        if let mir::StatementKind::StorageDead(local) = stmt.kind {
+            state.kill(local);
         }
     }
 
@@ -644,16 +642,16 @@ impl<'tcx> Analysis<'tcx> for EverInitializedPlaces<'_, 'tcx> {
     ) -> TerminatorEdges<'mir, 'tcx> {
         let move_data = self.move_data();
         let init_loc_map = &move_data.init_loc_map;
-        debug!(?terminator);
-        debug!("initializes move_indexes {:?}", init_loc_map[location]);
-        state.gen_all(
-            init_loc_map[location]
-                .iter()
-                .filter(|init_index| {
-                    move_data.inits[**init_index].kind != InitKind::NonPanicPathOnly
-                })
-                .copied(),
-        );
+
+        // Record inits of locals. Projections can be ignored.
+        state.gen_all(init_loc_map[location].iter().copied().filter_map(|ii| {
+            let init = &move_data.inits[ii];
+            if init.kind != InitKind::NonPanicPathOnly {
+                move_data.move_paths[init.path].place.as_local()
+            } else {
+                None
+            }
+        }));
         terminator.edges()
     }
 
@@ -666,9 +664,77 @@ impl<'tcx> Analysis<'tcx> for EverInitializedPlaces<'_, 'tcx> {
         let move_data = self.move_data();
         let init_loc_map = &move_data.init_loc_map;
 
+        // Record inits of locals. Projections can be ignored.
         let call_loc = self.body.terminator_loc(block);
-        for init_index in &init_loc_map[call_loc] {
-            state.gen_(*init_index);
+        state.gen_all(init_loc_map[call_loc].iter().copied().filter_map(|ii| {
+            let init = &move_data.inits[ii];
+            if init.kind == InitKind::NonPanicPathOnly {
+                move_data.move_paths[init.path].place.as_local()
+            } else {
+                None
+            }
+        }));
+    }
+}
+
+impl EverInitializedPlaces<'_, '_> {
+    /// Whether the init of `local` at `init` can reach `target` via a path that doesn't pass
+    /// through a `StorageDead(local)`. Mirrors the gen/kill structure of `EverInitializedPlaces`.
+    pub fn init_reaches_location(
+        body: &Body<'_>,
+        local: Local,
+        init: Init,
+        target: Location,
+    ) -> bool {
+        let init_loc = match init.location {
+            // Arguments are initialized on entry, and `StorageDead` is never emitted for them, so
+            // they reach every location.
+            InitLocation::Argument(_) => return true,
+            InitLocation::Statement(init_loc) => init_loc,
+        };
+
+        // Worklist of locations to walk forward from, seeded with the location(s) following `init`.
+        let mut queue = vec![];
+
+        let basic_blocks = &body.basic_blocks;
+        let init_block_data = &basic_blocks[init_loc.block];
+        if init_loc.statement_index < init_block_data.statements.len() {
+            // This case mirrors `apply_primary_statement_effect`.
+            queue.push(init_loc.successor_within_block());
+        } else if init.kind == InitKind::NonPanicPathOnly {
+            // This case mirrors `apply_call_return_effect`.
+            let TerminatorEdges::AssignOnReturn { return_, .. } =
+                init_block_data.terminator().edges()
+            else {
+                bug!("`NonPanicPathOnly` should only be seen on terminators with return edges");
+            };
+            queue.extend(return_.into_iter().map(BasicBlock::start_location));
+        } else {
+            // This case mirrors `apply_primary_terminator_effect`.
+            queue.extend(init_block_data.terminator().successors().map(BasicBlock::start_location));
+        }
+
+        let mut visited = FxIndexSet::default();
+        'outer: while let Some(loc) = queue.pop() {
+            if !visited.insert(loc) {
+                continue;
+            }
+            // Walk from `loc` to the end of its block, looking for `target` or a kill.
+            let block_data = &basic_blocks[loc.block];
+            for statement_index in loc.statement_index..=block_data.statements.len() {
+                if target == (Location { block: loc.block, statement_index }) {
+                    return true;
+                }
+                if let Some(stmt) = block_data.statements.get(statement_index)
+                    && let StatementKind::StorageDead(dead) = stmt.kind
+                    && dead == local
+                {
+                    continue 'outer;
+                }
+            }
+
+            queue.extend(block_data.terminator().successors().map(BasicBlock::start_location));
         }
+        false
     }
 }