Store locals in borrowck in BitSets

compiler/rustc_borrowck/src/diagnostics/explain_borrow.rs

@@ -630,7 +630,7 @@ impl<'tcx> MirBorrowckCtxt<'_, '_, 'tcx> {
         // ensure that we don't emit live boring locals as explanations.
         let is_local_boring = |local| {
             if let Some(polonius_context) = self.polonius_context {
-                polonius_context.boring_nll_locals.contains(&local)
+                polonius_context.boring_nll_locals.contains(local)
             } else {
                 assert!(!tcx.sess.opts.unstable_opts.polonius.is_next_enabled());
 

compiler/rustc_borrowck/src/lib.rs

@@ -29,7 +29,7 @@ use rustc_data_structures::graph::dominators::Dominators;
 use rustc_hir as hir;
 use rustc_hir::CRATE_HIR_ID;
 use rustc_hir::def_id::LocalDefId;
-use rustc_index::bit_set::MixedBitSet;
+use rustc_index::bit_set::{DenseBitSet, MixedBitSet};
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_infer::infer::outlives::env::RegionBoundPairs;
 use rustc_infer::infer::{
@@ -486,7 +486,7 @@ fn borrowck_check_region_constraints<'diag, 'tcx>(
             reservation_error_reported: Default::default(),
             uninitialized_error_reported: Default::default(),
             regioncx: &regioncx,
-            used_mut: Default::default(),
+            used_mut: DenseBitSet::new_empty(body.local_decls.len()),
             used_mut_upvars: SmallVec::new(),
             borrow_set: &borrow_set,
             upvars: &[],
@@ -525,7 +525,7 @@ fn borrowck_check_region_constraints<'diag, 'tcx>(
         reservation_error_reported: Default::default(),
         uninitialized_error_reported: Default::default(),
         regioncx: &regioncx,
-        used_mut: Default::default(),
+        used_mut: DenseBitSet::new_empty(body.local_decls.len()),
         used_mut_upvars: SmallVec::new(),
         borrow_set: &borrow_set,
         upvars: tcx.closure_captures(def),
@@ -560,17 +560,22 @@ fn borrowck_check_region_constraints<'diag, 'tcx>(
     // Note that this set is expected to be small - only upvars from closures
     // would have a chance of erroneously adding non-user-defined mutable vars
     // to the set.
-    let temporary_used_locals: FxIndexSet<Local> = mbcx
-        .used_mut
+    let mut temporary_used_locals = DenseBitSet::new_empty(body.local_decls.len());
+    mbcx.used_mut
         .iter()
-        .filter(|&local| !mbcx.body.local_decls[*local].is_user_variable())
-        .cloned()
-        .collect();
+        .filter(|local| !mbcx.body.local_decls[*local].is_user_variable())
+        .for_each(|local| {
+            temporary_used_locals.insert(local);
+        });
+
     // For the remaining unused locals that are marked as mutable, we avoid linting any that
     // were never initialized. These locals may have been removed as unreachable code; or will be
     // linted as unused variables.
-    let unused_mut_locals =
-        mbcx.body.mut_vars_iter().filter(|local| !mbcx.used_mut.contains(local)).collect();
+    let mut unused_mut_locals = DenseBitSet::new_empty(body.local_decls.len());
+    mbcx.body.mut_vars_iter().filter(|local| !mbcx.used_mut.contains(*local)).for_each(|local| {
+        unused_mut_locals.insert(local);
+    });
+
     mbcx.gather_used_muts(temporary_used_locals, unused_mut_locals);
 
     debug!("mbcx.used_mut: {:?}", mbcx.used_mut);
@@ -769,7 +774,7 @@ pub(crate) struct MirBorrowckCtxt<'a, 'diag, 'tcx> {
     uninitialized_error_reported: FxIndexSet<Local>,
     /// This field keeps track of all the local variables that are declared mut and are mutated.
     /// Used for the warning issued by an unused mutable local variable.
-    used_mut: FxIndexSet<Local>,
+    used_mut: DenseBitSet<Local>,
     /// If the function we're checking is a closure, then we'll need to report back the list of
     /// mutable upvars that have been used. This field keeps track of them.
     used_mut_upvars: SmallVec<[FieldIdx; 8]>,
@@ -2731,7 +2736,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
     fn lint_unused_mut(&self) {
         let tcx = self.infcx.tcx;
         let body = self.body;
-        for local in body.mut_vars_and_args_iter().filter(|local| !self.used_mut.contains(local)) {
+        for local in body.mut_vars_and_args_iter().filter(|local| !self.used_mut.contains(*local)) {
             let local_decl = &body.local_decls[local];
             let ClearCrossCrate::Set(SourceScopeLocalData { lint_root, .. }) =
                 body.source_scopes[local_decl.source_info.scope].local_data

compiler/rustc_borrowck/src/polonius/mod.rs

@@ -40,8 +40,7 @@ mod liveness_constraints;
 
 use std::collections::BTreeMap;
 
-use rustc_data_structures::fx::FxHashSet;
-use rustc_index::bit_set::SparseBitMatrix;
+use rustc_index::bit_set::{DenseBitSet, SparseBitMatrix};
 use rustc_middle::mir::{Body, Local};
 use rustc_middle::ty::RegionVid;
 use rustc_mir_dataflow::points::PointIndex;
@@ -58,7 +57,6 @@ pub(crate) type LiveLoans = SparseBitMatrix<PointIndex, BorrowIndex>;
 ///  - liveness data, created during MIR typeck, and which will be used to lazily compute the
 ///    polonius localized constraints, during NLL region inference as well as MIR dumping,
 ///  - data needed by the borrowck error computation and diagnostics.
-#[derive(Default)]
 pub(crate) struct PoloniusContext {
     /// The graph from which we extract the localized outlives constraints.
     graph: Option<LocalizedConstraintGraph>,
@@ -71,7 +69,16 @@ pub(crate) struct PoloniusContext {
     /// boring locals. A boring local is one whose type contains only such regions. Polonius
     /// currently has more boring locals than NLLs so we record the latter to use in errors and
     /// diagnostics, to focus on the locals we consider relevant and match NLL diagnostics.
-    pub(crate) boring_nll_locals: FxHashSet<Local>,
+    pub(crate) boring_nll_locals: DenseBitSet<Local>,
+}
+impl Default for PoloniusContext {
+    fn default() -> Self {
+        Self {
+            graph: None,
+            live_region_variances: Default::default(),
+            boring_nll_locals: DenseBitSet::new_empty(0),
+        }
+    }
 }
 
 /// The direction a constraint can flow into. Used to create liveness constraints according to

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

@@ -82,7 +82,7 @@ impl<'a> Iterator for AppearancesIter<'a> {
 
 impl LocalUseMap {
     pub(crate) fn build(
-        live_locals: &[Local],
+        live_locals: &DenseBitSet<Local>,
         location_map: &DenseLocationMap,
         body: &Body<'_>,
     ) -> Self {
@@ -98,11 +98,7 @@ impl LocalUseMap {
             return local_use_map;
         }
 
-        let mut locals_with_use_data: DenseBitSet<Local> =
-            DenseBitSet::new_empty(body.local_decls.len());
-        live_locals.iter().for_each(|&local| {
-            locals_with_use_data.insert(local);
-        });
+        let locals_with_use_data = live_locals;
 
         LocalUseMapBuild { local_use_map: &mut local_use_map, location_map, locals_with_use_data }
             .visit_body(body);
@@ -137,7 +133,7 @@ struct LocalUseMapBuild<'me> {
     // obtained the same information from `live_locals` but we want to
     // avoid repeatedly calling `Vec::contains()` (see `LocalUseMap` for
     // the rationale on the time-memory trade-off we're favoring here).
-    locals_with_use_data: DenseBitSet<Local>,
+    locals_with_use_data: &'me DenseBitSet<Local>,
 }
 
 impl Visitor<'_> for LocalUseMapBuild<'_> {

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

@@ -1,5 +1,5 @@
-use itertools::{Either, Itertools};
 use rustc_data_structures::fx::FxHashSet;
+use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::{TyContext, Visitor};
 use rustc_middle::mir::{Body, Local, Location, SourceInfo};
 use rustc_middle::span_bug;
@@ -55,8 +55,7 @@ pub(super) fn generate<'tcx>(
     {
         let (_, boring_locals) =
             compute_relevant_live_locals(typeck.tcx(), &free_regions, typeck.body);
-        typeck.polonius_context.as_mut().unwrap().boring_nll_locals =
-            boring_locals.into_iter().collect();
+        typeck.polonius_context.as_mut().unwrap().boring_nll_locals = boring_locals;
         free_regions = typeck.universal_regions.universal_regions_iter().collect();
     }
     let (relevant_live_locals, boring_locals) =
@@ -84,18 +83,20 @@ fn compute_relevant_live_locals<'tcx>(
     tcx: TyCtxt<'tcx>,
     free_regions: &FxHashSet<RegionVid>,
     body: &Body<'tcx>,
-) -> (Vec<Local>, Vec<Local>) {
-    let (boring_locals, relevant_live_locals): (Vec<_>, Vec<_>) =
-        body.local_decls.iter_enumerated().partition_map(|(local, local_decl)| {
-            if tcx.all_free_regions_meet(&local_decl.ty, |r| free_regions.contains(&r.as_var())) {
-                Either::Left(local)
-            } else {
-                Either::Right(local)
-            }
-        });
+) -> (DenseBitSet<Local>, DenseBitSet<Local>) {
+    let mut boring_locals = DenseBitSet::new_empty(body.local_decls.len());
+    let mut relevant_live_locals = DenseBitSet::new_empty(body.local_decls.len());
+
+    body.local_decls.iter_enumerated().for_each(|(local, local_decl)| {
+        if tcx.all_free_regions_meet(&local_decl.ty, |r| free_regions.contains(&r.as_var())) {
+            boring_locals.insert(local);
+        } else {
+            relevant_live_locals.insert(local);
+        }
+    });
 
     debug!("{} total variables", body.local_decls.len());
-    debug!("{} variables need liveness", relevant_live_locals.len());
+    debug!("{} variables need liveness", relevant_live_locals.count());
     debug!("{} regions outlive free regions", free_regions.len());
 
     (relevant_live_locals, boring_locals)

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

@@ -1,4 +1,4 @@
-use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_index::interval::IntervalSet;
 use rustc_infer::infer::canonical::QueryRegionConstraints;
@@ -42,8 +42,8 @@ pub(super) fn trace<'tcx>(
     typeck: &mut TypeChecker<'_, 'tcx>,
     location_map: &DenseLocationMap,
     move_data: &MoveData<'tcx>,
-    relevant_live_locals: Vec<Local>,
-    boring_locals: Vec<Local>,
+    relevant_live_locals: DenseBitSet<Local>,
+    boring_locals: DenseBitSet<Local>,
 ) {
     let _timer = typeck.tcx().prof.generic_activity("borrowck_liveness_trace");
 
@@ -131,8 +131,8 @@ impl<'a, 'typeck, 'tcx> LivenessResults<'a, 'typeck, 'tcx> {
         }
     }
 
-    fn compute_for_all_locals(&mut self, relevant_live_locals: Vec<Local>) {
-        for local in relevant_live_locals {
+    fn compute_for_all_locals(&mut self, relevant_live_locals: DenseBitSet<Local>) {
+        for local in relevant_live_locals.iter() {
             self.reset_local_state();
             self.add_defs_for(local);
             self.compute_use_live_points_for(local);
@@ -161,8 +161,8 @@ impl<'a, 'typeck, 'tcx> LivenessResults<'a, 'typeck, 'tcx> {
     /// These are all the locals which do not potentially reference a region local
     /// to this body. Locals which only reference free regions are always drop-live
     /// and can therefore safely be dropped.
-    fn dropck_boring_locals(&mut self, boring_locals: Vec<Local>) {
-        for local in boring_locals {
+    fn dropck_boring_locals(&mut self, boring_locals: DenseBitSet<Local>) {
+        for local in boring_locals.iter() {
             let local_ty = self.cx.body().local_decls[local].ty;
             let local_span = self.cx.body().local_decls[local].source_info.span;
             let drop_data = self.cx.drop_data.entry(local_ty).or_insert_with({
@@ -182,7 +182,7 @@ impl<'a, 'typeck, 'tcx> LivenessResults<'a, 'typeck, 'tcx> {
     ///
     /// Add facts for all locals with free regions, since regions may outlive
     /// the function body only at certain nodes in the CFG.
-    fn add_extra_drop_facts(&mut self, relevant_live_locals: &[Local]) {
+    fn add_extra_drop_facts(&mut self, relevant_live_locals: &DenseBitSet<Local>) {
         // This collect is more necessary than immediately apparent
         // because these facts go into `add_drop_live_facts_for()`,
         // which also writes to `polonius_facts`, and so this is genuinely
@@ -194,15 +194,12 @@ impl<'a, 'typeck, 'tcx> LivenessResults<'a, 'typeck, 'tcx> {
         // `add_drop_live_facts_for()` that make sense.
         let Some(facts) = self.cx.typeck.polonius_facts.as_ref() else { return };
         let facts_to_add: Vec<_> = {
-            let relevant_live_locals: FxIndexSet<_> =
-                relevant_live_locals.iter().copied().collect();
-
             facts
                 .var_dropped_at
                 .iter()
                 .filter_map(|&(local, location_index)| {
                     let local_ty = self.cx.body().local_decls[local].ty;
-                    if relevant_live_locals.contains(&local) || !local_ty.has_free_regions() {
+                    if relevant_live_locals.contains(local) || !local_ty.has_free_regions() {
                         return None;
                     }
 

compiler/rustc_borrowck/src/used_muts.rs

@@ -1,4 +1,4 @@
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::{PlaceContext, Visitor};
 use rustc_middle::mir::{
     Local, Location, Place, Statement, StatementKind, Terminator, TerminatorKind,
@@ -25,8 +25,8 @@ impl<'tcx> MirBorrowckCtxt<'_, '_, 'tcx> {
     ///  See #55344 for context.
     pub(crate) fn gather_used_muts(
         &mut self,
-        temporary_used_locals: FxIndexSet<Local>,
-        mut never_initialized_mut_locals: FxIndexSet<Local>,
+        temporary_used_locals: DenseBitSet<Local>,
+        mut never_initialized_mut_locals: DenseBitSet<Local>,
     ) {
         {
             let mut visitor = GatherUsedMutsVisitor {
@@ -40,15 +40,15 @@ impl<'tcx> MirBorrowckCtxt<'_, '_, 'tcx> {
         // Take the union of the existed `used_mut` set with those variables we've found were
         // never initialized.
         debug!("gather_used_muts: never_initialized_mut_locals={:?}", never_initialized_mut_locals);
-        self.used_mut = self.used_mut.union(&never_initialized_mut_locals).cloned().collect();
+        let _ = self.used_mut.union(&never_initialized_mut_locals);
     }
 }
 
 /// MIR visitor for collecting used mutable variables.
 /// The 'visit lifetime represents the duration of the MIR walk.
 struct GatherUsedMutsVisitor<'a, 'b, 'diag, 'tcx> {
-    temporary_used_locals: FxIndexSet<Local>,
-    never_initialized_mut_locals: &'a mut FxIndexSet<Local>,
+    temporary_used_locals: DenseBitSet<Local>,
+    never_initialized_mut_locals: &'a mut DenseBitSet<Local>,
     mbcx: &'a mut MirBorrowckCtxt<'b, 'diag, 'tcx>,
 }
 
@@ -59,8 +59,7 @@ impl GatherUsedMutsVisitor<'_, '_, '_, '_> {
         // be those that were never initialized - we will consider those as being used as
         // they will either have been removed by unreachable code optimizations; or linted
         // as unused variables.
-        // FIXME(#120456) - is `swap_remove` correct?
-        self.never_initialized_mut_locals.swap_remove(&into.local);
+        self.never_initialized_mut_locals.remove(into.local);
     }
 }
 
@@ -91,7 +90,7 @@ impl<'tcx> Visitor<'tcx> for GatherUsedMutsVisitor<'_, '_, '_, 'tcx> {
     }
 
     fn visit_local(&mut self, local: Local, place_context: PlaceContext, location: Location) {
-        if place_context.is_place_assignment() && self.temporary_used_locals.contains(&local) {
+        if place_context.is_place_assignment() && self.temporary_used_locals.contains(local) {
             // Propagate the Local assigned at this Location as a used mutable local variable
             for moi in &self.mbcx.move_data.move_out_loc_map[location] {
                 let mpi = &self.mbcx.move_data.move_outs[*moi].path;