Simplify MaybeRequiresStorage

compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs

@@ -1,12 +1,13 @@
 use std::borrow::Cow;
-use std::cell::RefCell;
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::visit::{NonMutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
+use smallvec::SmallVec;
 
 use super::MaybeBorrowedLocals;
-use crate::{Analysis, GenKill, ResultsCursor};
+use crate::{Analysis, GenKill, Results, ResultsCursor};
 
 /// The set of locals in a MIR body that do not have `StorageLive`/`StorageDead` annotations.
 ///
@@ -111,21 +112,51 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeStorageDead<'a> {
     }
 }
 
-type BorrowedLocalsResults<'mir, 'tcx> = ResultsCursor<'mir, 'tcx, MaybeBorrowedLocals>;
+/// For each location, records which locals can be killed by `MaybeRequiresStorage`.
+type KillableLocals = FxHashMap<Location, SmallVec<[Local; 4]>>;
 
 /// Dataflow analysis that determines whether each local requires storage at a
 /// given location; i.e. whether its storage can go away without being observed.
-pub struct MaybeRequiresStorage<'mir, 'tcx> {
-    borrowed_locals: RefCell<BorrowedLocalsResults<'mir, 'tcx>>,
+pub struct MaybeRequiresStorage {
+    /// Used to kill locals that are fully moved and have not been borrowed.
+    killable_locals: KillableLocals,
 }
 
-impl<'mir, 'tcx> MaybeRequiresStorage<'mir, 'tcx> {
-    pub fn new(borrowed_locals: BorrowedLocalsResults<'mir, 'tcx>) -> Self {
-        MaybeRequiresStorage { borrowed_locals: RefCell::new(borrowed_locals) }
+impl MaybeRequiresStorage {
+    pub fn new<'tcx>(
+        body: &Body<'tcx>,
+        borrowed_locals: &Results<'tcx, MaybeBorrowedLocals>,
+    ) -> Self {
+        struct KillableLocalsVisitor<'mir, 'tcx> {
+            borrowed_locals_cursor: ResultsCursor<'mir, 'tcx, MaybeBorrowedLocals>,
+            killable_locals: KillableLocals,
+        }
+
+        impl<'tcx> Visitor<'tcx> for KillableLocalsVisitor<'_, 'tcx> {
+            fn visit_local(&mut self, local: Local, context: PlaceContext, loc: Location) {
+                if PlaceContext::NonMutatingUse(NonMutatingUseContext::Move) == context {
+                    self.borrowed_locals_cursor.seek_before_primary_effect(loc);
+                    if !self.borrowed_locals_cursor.get().contains(local) {
+                        self.killable_locals.entry(loc).or_default().push(local);
+                    }
+                }
+            }
+        }
+
+        let mut visitor = KillableLocalsVisitor {
+            borrowed_locals_cursor: ResultsCursor::new_borrowing(body, borrowed_locals),
+            killable_locals: Default::default(),
+        };
+
+        for (bb, data) in body.basic_blocks.iter_enumerated() {
+            visitor.visit_basic_block_data(bb, data);
+        }
+
+        MaybeRequiresStorage { killable_locals: visitor.killable_locals }
     }
 }
 
-impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage<'_, 'tcx> {
+impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage {
     type Domain = DenseBitSet<Local>;
 
     const NAME: &'static str = "requires_storage";
@@ -182,8 +213,7 @@ impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage<'_, 'tcx> {
         stmt: &Statement<'tcx>,
         loc: Location,
     ) {
-        // If we move from a place then it only stops needing storage *after*
-        // that statement.
+        // If we move from a place then it only stops needing storage *after* that statement.
         self.check_for_move(state, loc);
 
         match &stmt.kind {
@@ -316,27 +346,12 @@ impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage<'_, 'tcx> {
     }
 }
 
-impl<'tcx> MaybeRequiresStorage<'_, 'tcx> {
+impl MaybeRequiresStorage {
     /// Kill locals that are fully moved and have not been borrowed.
-    fn check_for_move(&self, state: &mut <Self as Analysis<'tcx>>::Domain, loc: Location) {
-        let mut borrowed_locals = self.borrowed_locals.borrow_mut();
-        let body = borrowed_locals.body();
-        let mut visitor = MoveVisitor { state, borrowed_locals: &mut borrowed_locals };
-        visitor.visit_location(body, loc);
-    }
-}
-
-struct MoveVisitor<'a, 'mir, 'tcx> {
-    borrowed_locals: &'a mut BorrowedLocalsResults<'mir, 'tcx>,
-    state: &'a mut DenseBitSet<Local>,
-}
-
-impl<'tcx> Visitor<'tcx> for MoveVisitor<'_, '_, 'tcx> {
-    fn visit_local(&mut self, local: Local, context: PlaceContext, loc: Location) {
-        if PlaceContext::NonMutatingUse(NonMutatingUseContext::Move) == context {
-            self.borrowed_locals.seek_before_primary_effect(loc);
-            if !self.borrowed_locals.get().contains(local) {
-                self.state.kill(local);
+    fn check_for_move(&self, state: &mut <Self as Analysis<'_>>::Domain, loc: Location) {
+        if let Some(locals) = self.killable_locals.get(&loc) {
+            for &l in locals {
+                state.kill(l);
             }
         }
     }

compiler/rustc_mir_transform/src/coroutine/layout.rs

@@ -93,12 +93,10 @@ pub(super) fn locals_live_across_suspend_points<'tcx>(
 
     // Calculate the MIR locals that have been previously borrowed (even if they are still active).
     let borrowed_locals = MaybeBorrowedLocals.iterate_to_fixpoint(tcx, body, Some("coroutine"));
-    let borrowed_locals_cursor1 = ResultsCursor::new_borrowing(body, &borrowed_locals);
-    let mut borrowed_locals_cursor2 = ResultsCursor::new_borrowing(body, &borrowed_locals);
 
     // Calculate the MIR locals that we need to keep storage around for.
     let requires_storage =
-        MaybeRequiresStorage::new(borrowed_locals_cursor1).iterate_to_fixpoint(tcx, body, None);
+        MaybeRequiresStorage::new(body, &borrowed_locals).iterate_to_fixpoint(tcx, body, None);
     let mut requires_storage_cursor = ResultsCursor::new_borrowing(body, &requires_storage);
 
     // Calculate the liveness of MIR locals ignoring borrows.
@@ -109,6 +107,7 @@ pub(super) fn locals_live_across_suspend_points<'tcx>(
     let mut live_locals_at_suspension_points = Vec::new();
     let mut source_info_at_suspension_points = Vec::new();
     let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
+    let mut borrowed_locals_cursor = ResultsCursor::new_owning(body, borrowed_locals);
 
     for (block, data) in body.basic_blocks.iter_enumerated() {
         let TerminatorKind::Yield { .. } = data.terminator().kind else { continue };
@@ -129,8 +128,8 @@ pub(super) fn locals_live_across_suspend_points<'tcx>(
             // If a borrow is converted to a raw reference, we must also assume that it lives
             // forever. Note that the final liveness is still bounded by the storage liveness
             // of the local, which happens using the `intersect` operation below.
-            borrowed_locals_cursor2.seek_before_primary_effect(loc);
-            live_locals.union(borrowed_locals_cursor2.get());
+            borrowed_locals_cursor.seek_before_primary_effect(loc);
+            live_locals.union(borrowed_locals_cursor.get());
         }
 
         // Store the storage liveness for later use so we can restore the state
@@ -236,7 +235,7 @@ fn compute_storage_conflicts<'mir, 'tcx>(
     body: &'mir Body<'tcx>,
     saved_locals: &'mir CoroutineSavedLocals,
     always_live_locals: DenseBitSet<Local>,
-    results: &Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>,
+    results: &Results<'tcx, MaybeRequiresStorage>,
 ) -> BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal> {
     assert_eq!(body.local_decls.len(), saved_locals.domain_size());
 
@@ -294,12 +293,10 @@ struct StorageConflictVisitor<'a, 'tcx> {
     eligible_storage_live: DenseBitSet<Local>,
 }
 
-impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage<'a, 'tcx>>
-    for StorageConflictVisitor<'a, 'tcx>
-{
+impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage> for StorageConflictVisitor<'a, 'tcx> {
     fn visit_after_early_statement_effect(
         &mut self,
-        _analysis: &MaybeRequiresStorage<'a, 'tcx>,
+        _analysis: &MaybeRequiresStorage,
         state: &DenseBitSet<Local>,
         _statement: &Statement<'tcx>,
         loc: Location,
@@ -309,7 +306,7 @@ impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage<'a, 'tcx>>
 
     fn visit_after_early_terminator_effect(
         &mut self,
-        _analysis: &MaybeRequiresStorage<'a, 'tcx>,
+        _analysis: &MaybeRequiresStorage,
         state: &DenseBitSet<Local>,
         _terminator: &Terminator<'tcx>,
         loc: Location,