Fix liveness analysis for yield terminators

A resume place is evaluated and assigned to only after a yield
terminator resumes. Ensure that locals used when evaluating the
resume place are live across the yield.

Author: tmiasko

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -23,12 +23,6 @@ use crate::{Analysis, AnalysisDomain, Backward, CallReturnPlaces, GenKill, GenKi
 /// [liveness]: https://en.wikipedia.org/wiki/Live_variable_analysis
 pub struct MaybeLiveLocals;
 
-impl MaybeLiveLocals {
-    fn transfer_function<'a, T>(&self, trans: &'a mut T) -> TransferFunction<'a, T> {
-        TransferFunction(trans)
-    }
-}
-
 impl<'tcx> AnalysisDomain<'tcx> for MaybeLiveLocals {
     type Domain = ChunkedBitSet<Local>;
     type Direction = Backward;
@@ -54,7 +48,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeLiveLocals {
         statement: &mir::Statement<'tcx>,
         location: Location,
     ) {
-        self.transfer_function(trans).visit_statement(statement, location);
+        TransferFunction(trans).visit_statement(statement, location);
     }
 
     fn terminator_effect(
@@ -63,7 +57,7 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeLiveLocals {
         terminator: &mir::Terminator<'tcx>,
         location: Location,
     ) {
-        self.transfer_function(trans).visit_terminator(terminator, location);
+        TransferFunction(trans).visit_terminator(terminator, location);
     }
 
     fn call_return_effect(
@@ -85,9 +79,11 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeLiveLocals {
         _resume_block: mir::BasicBlock,
         resume_place: mir::Place<'tcx>,
     ) {
-        if let Some(local) = resume_place.as_local() {
-            trans.kill(local);
-        }
+        YieldResumeEffect(trans).visit_place(
+            &resume_place,
+            PlaceContext::MutatingUse(MutatingUseContext::Yield),
+            Location::START,
+        )
     }
 }
 
@@ -98,23 +94,58 @@ where
     T: GenKill<Local>,
 {
     fn visit_place(&mut self, place: &mir::Place<'tcx>, context: PlaceContext, location: Location) {
-        let local = place.local;
+        if let PlaceContext::MutatingUse(MutatingUseContext::Yield) = context {
+            // The resume place is evaluated and assigned to only after generator resumes, so its
+            // effect is handled separately in `yield_resume_effect`.
+            return;
+        }
+
+        match DefUse::for_place(*place, context) {
+            Some(DefUse::Def) => {
+                if let PlaceContext::MutatingUse(
+                    MutatingUseContext::Call | MutatingUseContext::AsmOutput,
+                ) = context
+                {
+                    // For the associated terminators, this is only a `Def` when the terminator returns
+                    // "successfully." As such, we handle this case separately in `call_return_effect`
+                    // above. However, if the place looks like `*_5`, this is still unconditionally a use of
+                    // `_5`.
+                } else {
+                    self.0.kill(place.local);
+                }
+            }
+            Some(DefUse::Use) => self.0.gen(place.local),
+            None => {}
+        }
 
-        // We purposefully do not call `super_place` here to avoid calling `visit_local` for this
-        // place with one of the `Projection` variants of `PlaceContext`.
         self.visit_projection(place.as_ref(), context, location);
+    }
 
-        match DefUse::for_place(*place, context) {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, _: Location) {
+        match DefUse::for_place(local.into(), context) {
             Some(DefUse::Def) => self.0.kill(local),
             Some(DefUse::Use) => self.0.gen(local),
             None => {}
         }
     }
+}
+
+struct YieldResumeEffect<'a, T>(&'a mut T);
+
+impl<'tcx, T> Visitor<'tcx> for YieldResumeEffect<'_, T>
+where
+    T: GenKill<Local>,
+{
+    fn visit_place(&mut self, place: &mir::Place<'tcx>, context: PlaceContext, location: Location) {
+        match DefUse::for_place(*place, context) {
+            Some(DefUse::Def) => self.0.kill(place.local),
+            Some(DefUse::Use) => self.0.gen(place.local),
+            None => {}
+        }
+        self.visit_projection(place.as_ref(), context, location);
+    }
 
     fn visit_local(&mut self, local: Local, context: PlaceContext, _: Location) {
-        // Because we do not call `super_place` above, `visit_local` is only called for locals that
-        // do not appear as part of  a `Place` in the MIR. This handles cases like the implicit use
-        // of the return place in a `Return` terminator or the index in an `Index` projection.
         match DefUse::for_place(local.into(), context) {
             Some(DefUse::Def) => self.0.kill(local),
             Some(DefUse::Use) => self.0.gen(local),
@@ -134,7 +165,13 @@ impl DefUse {
         match context {
             PlaceContext::NonUse(_) => None,
 
-            PlaceContext::MutatingUse(MutatingUseContext::Store | MutatingUseContext::Deinit) => {
+            PlaceContext::MutatingUse(
+                MutatingUseContext::Call
+                | MutatingUseContext::Yield
+                | MutatingUseContext::AsmOutput
+                | MutatingUseContext::Store
+                | MutatingUseContext::Deinit,
+            ) => {
                 if place.is_indirect() {
                     // Treat derefs as a use of the base local. `*p = 4` is not a def of `p` but a
                     // use.
@@ -152,16 +189,6 @@ impl DefUse {
                 place.is_indirect().then_some(DefUse::Use)
             }
 
-            // For the associated terminators, this is only a `Def` when the terminator returns
-            // "successfully." As such, we handle this case separately in `call_return_effect`
-            // above. However, if the place looks like `*_5`, this is still unconditionally a use of
-            // `_5`.
-            PlaceContext::MutatingUse(
-                MutatingUseContext::Call
-                | MutatingUseContext::Yield
-                | MutatingUseContext::AsmOutput,
-            ) => place.is_indirect().then_some(DefUse::Use),
-
             // All other contexts are uses...
             PlaceContext::MutatingUse(
                 MutatingUseContext::AddressOf
@@ -290,8 +317,10 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
         _resume_block: mir::BasicBlock,
         resume_place: mir::Place<'tcx>,
     ) {
-        if let Some(local) = resume_place.as_local() {
-            trans.remove(local);
-        }
+        YieldResumeEffect(trans).visit_place(
+            &resume_place,
+            PlaceContext::MutatingUse(MutatingUseContext::Yield),
+            Location::START,
+        )
     }
 }