Auto merge of #81132 - bugadani:map-prealloc, r=matthewjasper

Borrowck: refactor visited map to a bitset

This PR refactors `Borrows`  and the `precompute_borrows_out_of_scope` function so that this initial phase has a much reduced memory pressure. This is achieved by reducing what is stored on the heap, and also reusing heap memory as much as possible.

Cargo.lock

@@ -5395,7 +5395,7 @@ dependencies = [
  "chrono",
  "lazy_static",
  "matchers",
- "parking_lot 0.11.0",
+ "parking_lot 0.9.0",
  "regex",
  "serde",
  "serde_json",

compiler/rustc_mir/src/borrow_check/mod.rs

@@ -243,7 +243,7 @@ fn do_mir_borrowck<'a, 'tcx>(
 
     let regioncx = Rc::new(regioncx);
 
-    let flow_borrows = Borrows::new(tcx, &body, regioncx.clone(), &borrow_set)
+    let flow_borrows = Borrows::new(tcx, &body, &regioncx, &borrow_set)
         .into_engine(tcx, &body)
         .pass_name("borrowck")
         .iterate_to_fixpoint();
@@ -287,7 +287,7 @@ fn do_mir_borrowck<'a, 'tcx>(
                 regioncx: regioncx.clone(),
                 used_mut: Default::default(),
                 used_mut_upvars: SmallVec::new(),
-                borrow_set: borrow_set.clone(),
+                borrow_set: Rc::clone(&borrow_set),
                 dominators,
                 upvars: Vec::new(),
                 local_names: IndexVec::from_elem(None, &promoted_body.local_decls),
@@ -317,10 +317,10 @@ fn do_mir_borrowck<'a, 'tcx>(
         move_error_reported: BTreeMap::new(),
         uninitialized_error_reported: Default::default(),
         errors_buffer,
-        regioncx,
+        regioncx: Rc::clone(&regioncx),
         used_mut: Default::default(),
         used_mut_upvars: SmallVec::new(),
-        borrow_set,
+        borrow_set: Rc::clone(&borrow_set),
         dominators,
         upvars,
         local_names,

compiler/rustc_mir/src/dataflow/impls/borrows.rs

@@ -11,7 +11,6 @@ use crate::borrow_check::{
 use crate::dataflow::{self, fmt::DebugWithContext, GenKill};
 
 use std::fmt;
-use std::rc::Rc;
 
 rustc_index::newtype_index! {
     pub struct BorrowIndex {
@@ -30,24 +29,39 @@ pub struct Borrows<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a Body<'tcx>,
 
-    borrow_set: Rc<BorrowSet<'tcx>>,
+    borrow_set: &'a BorrowSet<'tcx>,
     borrows_out_of_scope_at_location: FxHashMap<Location, Vec<BorrowIndex>>,
-
-    /// NLL region inference context with which NLL queries should be resolved
-    _nonlexical_regioncx: Rc<RegionInferenceContext<'tcx>>,
 }
 
 struct StackEntry {
     bb: mir::BasicBlock,
     lo: usize,
     hi: usize,
-    first_part_only: bool,
 }
 
-fn precompute_borrows_out_of_scope<'tcx>(
+struct OutOfScopePrecomputer<'a, 'tcx> {
-    body: &Body<'tcx>,
+    visited: BitSet<mir::BasicBlock>,
-    regioncx: &Rc<RegionInferenceContext<'tcx>>,
+    visit_stack: Vec<StackEntry>,
-    borrows_out_of_scope_at_location: &mut FxHashMap<Location, Vec<BorrowIndex>>,
+    body: &'a Body<'tcx>,
+    regioncx: &'a RegionInferenceContext<'tcx>,
+    borrows_out_of_scope_at_location: FxHashMap<Location, Vec<BorrowIndex>>,
+}
+
+impl<'a, 'tcx> OutOfScopePrecomputer<'a, 'tcx> {
+    fn new(body: &'a Body<'tcx>, regioncx: &'a RegionInferenceContext<'tcx>) -> Self {
+        OutOfScopePrecomputer {
+            visited: BitSet::new_empty(body.basic_blocks().len()),
+            visit_stack: vec![],
+            body,
+            regioncx,
+            borrows_out_of_scope_at_location: FxHashMap::default(),
+        }
+    }
+}
+
+impl<'tcx> OutOfScopePrecomputer<'_, 'tcx> {
+    fn precompute_borrows_out_of_scope(
+        &mut self,
         borrow_index: BorrowIndex,
         borrow_region: RegionVid,
         location: Location,
@@ -62,26 +76,27 @@ fn precompute_borrows_out_of_scope<'tcx>(
         // `visited` once they are added to `stack`, before they are actually
         // processed, because this avoids the need to look them up again on
         // completion.
-    let mut visited = FxHashMap::default();
+        self.visited.insert(location.block);
-    visited.insert(location.block, location.statement_index);
 
-    let mut stack = vec![];
+        let mut first_lo = location.statement_index;
-    stack.push(StackEntry {
+        let first_hi = self.body[location.block].statements.len();
-        bb: location.block,
-        lo: location.statement_index,
-        hi: body[location.block].statements.len(),
-        first_part_only: false,
-    });
 
-    while let Some(StackEntry { bb, lo, hi, first_part_only }) = stack.pop() {
+        self.visit_stack.push(StackEntry { bb: location.block, lo: first_lo, hi: first_hi });
-        let mut finished_early = first_part_only;
+
+        while let Some(StackEntry { bb, lo, hi }) = self.visit_stack.pop() {
+            // If we process the first part of the first basic block (i.e. we encounter that block
+            // for the second time), we no longer have to visit its successors again.
+            let mut finished_early = bb == location.block && hi != first_hi;
             for i in lo..=hi {
                 let location = Location { block: bb, statement_index: i };
                 // If region does not contain a point at the location, then add to list and skip
                 // successor locations.
-            if !regioncx.region_contains(borrow_region, location) {
+                if !self.regioncx.region_contains(borrow_region, location) {
                     debug!("borrow {:?} gets killed at {:?}", borrow_index, location);
-                borrows_out_of_scope_at_location.entry(location).or_default().push(borrow_index);
+                    self.borrows_out_of_scope_at_location
+                        .entry(location)
+                        .or_default()
+                        .push(borrow_index);
                     finished_early = true;
                     break;
                 }
@@ -89,73 +104,61 @@ fn precompute_borrows_out_of_scope<'tcx>(
 
             if !finished_early {
                 // Add successor BBs to the work list, if necessary.
-            let bb_data = &body[bb];
+                let bb_data = &self.body[bb];
-            assert!(hi == bb_data.statements.len());
+                debug_assert!(hi == bb_data.statements.len());
                 for &succ_bb in bb_data.terminator().successors() {
-                visited
+                    if self.visited.insert(succ_bb) == false {
-                    .entry(succ_bb)
+                        if succ_bb == location.block && first_lo > 0 {
-                    .and_modify(|lo| {
                             // `succ_bb` has been seen before. If it wasn't
                             // fully processed, add its first part to `stack`
                             // for processing.
-                        if *lo > 0 {
+                            self.visit_stack.push(StackEntry {
-                            stack.push(StackEntry {
                                 bb: succ_bb,
                                 lo: 0,
-                                hi: *lo - 1,
+                                hi: first_lo - 1,
-                                first_part_only: true,
                             });
-                        }
+
                             // And update this entry with 0, to represent the
                             // whole BB being processed.
-                        *lo = 0;
+                            first_lo = 0;
-                    })
+                        }
-                    .or_insert_with(|| {
+                    } else {
                         // succ_bb hasn't been seen before. Add it to
                         // `stack` for processing.
-                        stack.push(StackEntry {
+                        self.visit_stack.push(StackEntry {
                             bb: succ_bb,
                             lo: 0,
-                            hi: body[succ_bb].statements.len(),
+                            hi: self.body[succ_bb].statements.len(),
-                            first_part_only: false,
-                        });
-                        // Insert 0 for this BB, to represent the whole BB
-                        // being processed.
-                        0
                         });
                     }
                 }
             }
         }
 
+        self.visited.clear();
+    }
+}
+
 impl<'a, 'tcx> Borrows<'a, 'tcx> {
     crate fn new(
         tcx: TyCtxt<'tcx>,
         body: &'a Body<'tcx>,
-        nonlexical_regioncx: Rc<RegionInferenceContext<'tcx>>,
+        nonlexical_regioncx: &'a RegionInferenceContext<'tcx>,
-        borrow_set: &Rc<BorrowSet<'tcx>>,
+        borrow_set: &'a BorrowSet<'tcx>,
     ) -> Self {
-        let mut borrows_out_of_scope_at_location = FxHashMap::default();
+        let mut prec = OutOfScopePrecomputer::new(body, nonlexical_regioncx);
         for (borrow_index, borrow_data) in borrow_set.iter_enumerated() {
             let borrow_region = borrow_data.region.to_region_vid();
             let location = borrow_data.reserve_location;
 
-            precompute_borrows_out_of_scope(
+            prec.precompute_borrows_out_of_scope(borrow_index, borrow_region, location);
-                body,
-                &nonlexical_regioncx,
-                &mut borrows_out_of_scope_at_location,
-                borrow_index,
-                borrow_region,
-                location,
-            );
         }
 
         Borrows {
             tcx,
             body,
-            borrow_set: borrow_set.clone(),
+            borrow_set,
-            borrows_out_of_scope_at_location,
+            borrows_out_of_scope_at_location: prec.borrows_out_of_scope_at_location,
-            _nonlexical_regioncx: nonlexical_regioncx,
         }
     }