Properly compute dominance property for calls and yields.

Author: cjgillot

compiler/rustc_mir_transform/src/ssa.rs

@@ -50,9 +50,13 @@ impl SmallDominators<'_> {
         let assign_dominates = match *set {
             Set1::Empty | Set1::Many => false,
             Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(assign)) => {
+            Set1::One(LocationExtended::Assign(assign)) => {
                 self.dominates(assign.successor_within_block(), loc)
             }
+            Set1::One(LocationExtended::Terminator(_, effect_bb)) => {
+                let effect_loc = Location { block: effect_bb, statement_index: 0 };
+                self.dominates(effect_loc, loc)
+            }
         };
         // We are visiting a use that is not dominated by an assignment.
         // Either there is a cycle involved, or we are reading for uninitialized local.
@@ -73,7 +77,8 @@ impl SsaLocals {
         let dominators = SmallDominators { inner: dominators };
 
         let direct_uses = IndexVec::from_elem(0, &body.local_decls);
-        let mut visitor = SsaVisitor { assignments, assignment_order, dominators, direct_uses };
+        let mut visitor =
+            SsaVisitor { body, assignments, assignment_order, dominators, direct_uses };
 
         for local in body.args_iter() {
             visitor.assignments[local] = Set1::One(LocationExtended::Arg);
@@ -134,12 +139,9 @@ impl SsaLocals {
         body: &'a Body<'tcx>,
     ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a {
         self.assignment_order.iter().filter_map(|&local| {
-            if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local]
-                // `loc` must point to:
-                // - a direct assignment to `local`;
-                // - a call or a yield terminator.
-                && let Either::Left(stmt) = body.stmt_at(loc)
-            {
+            if let Set1::One(LocationExtended::Assign(loc)) = self.assignments[local] {
+                // `loc` must point to a direct assignment to `local`.
+                let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
                 assert_eq!(target.as_local(), Some(local));
                 Some((local, rvalue, loc))
@@ -197,18 +199,24 @@ impl SsaLocals {
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 enum LocationExtended {
-    Plain(Location),
+    /// This local was assigned as a function parameter, before any execution.
     Arg,
+    /// `Location` points to the `Assign` statement.
+    Assign(Location),
+    /// The blocks are respectively the bb which contains the assignment,
+    /// and the bb in which the assignment effect is complete.
+    Terminator(BasicBlock, BasicBlock),
 }
 
-struct SsaVisitor<'a> {
+struct SsaVisitor<'a, 'tcx> {
+    body: &'a Body<'tcx>,
     dominators: SmallDominators<'a>,
     assignments: IndexVec<Local, Set1<LocationExtended>>,
     assignment_order: Vec<Local>,
     direct_uses: IndexVec<Local, u32>,
 }
 
-impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
+impl<'tcx> Visitor<'tcx> for SsaVisitor<'_, 'tcx> {
     fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
         match ctxt {
             PlaceContext::MutatingUse(MutatingUseContext::Projection)
@@ -233,10 +241,20 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
     }
 
     fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) {
-        if let PlaceContext::MutatingUse(MutatingUseContext::Store | MutatingUseContext::Call | MutatingUseContext::Yield) = ctxt
+        let location = match ctxt {
+            PlaceContext::MutatingUse(MutatingUseContext::Store) => {
+                Some(LocationExtended::Assign(loc))
+            }
+            PlaceContext::MutatingUse(MutatingUseContext::Call | MutatingUseContext::Yield) => {
+                let target = self.body.basic_blocks[loc.block].terminator().successors().next();
+                target.map(|target| LocationExtended::Terminator(loc.block, target))
+            }
+            _ => None,
+        };
+        if let Some(location) = location
             && let Some(local) = place.as_local()
         {
-            self.assignments[local].insert(LocationExtended::Plain(loc));
+            self.assignments[local].insert(location);
             if let Set1::One(_) = self.assignments[local] {
                 // Only record if SSA-like, to avoid growing the vector needlessly.
                 self.assignment_order.push(local);
@@ -337,7 +355,7 @@ impl StorageLiveLocals {
             for (statement_index, statement) in bbdata.statements.iter().enumerate() {
                 if let StatementKind::StorageLive(local) = statement.kind {
                     storage_live[local]
-                        .insert(LocationExtended::Plain(Location { block, statement_index }));
+                        .insert(LocationExtended::Assign(Location { block, statement_index }));
                 }
             }
         }