Only consider a local to be SSA if assignment dominates all uses.

Author: cjgillot

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -1,4 +1,5 @@
 use either::Either;
+use rustc_data_structures::graph::dominators::Dominators;
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::IndexVec;
 use rustc_middle::middle::resolve_lifetime::Set1;
@@ -65,6 +66,7 @@ enum LocationExtended {
 
 #[derive(Debug)]
 struct SsaLocals {
+    dominators: Dominators<BasicBlock>,
     /// Assignments to each local.  This defines whether the local is SSA.
     assignments: IndexVec<Local, Set1<LocationExtended>>,
     /// We visit the body in reverse postorder, to ensure each local is assigned before it is used.
@@ -78,7 +80,8 @@ impl SsaLocals {
         let assignment_order = Vec::new();
 
         let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
-        let mut this = SsaLocals { assignments, assignment_order };
+        let dominators = body.basic_blocks.dominators();
+        let mut this = SsaLocals { assignments, assignment_order, dominators };
 
         let borrowed = borrowed_locals(body);
         for (local, decl) in body.local_decls.iter_enumerated() {
@@ -117,7 +120,23 @@ impl<'tcx> Visitor<'tcx> for SsaLocals {
             PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many,
             // Immutable borrows and AddressOf are taken into account in `SsaLocals::new` by
             // removing non-freeze locals.
-            PlaceContext::NonMutatingUse(_) | PlaceContext::NonUse(_) => {}
+            PlaceContext::NonMutatingUse(_) => {
+                let set = &mut self.assignments[local];
+                let assign_dominates = match *set {
+                    Set1::Empty | Set1::Many => false,
+                    Set1::One(LocationExtended::Arg) => true,
+                    Set1::One(LocationExtended::Plain(assign)) => {
+                        assign.dominates(loc, &self.dominators)
+                    }
+                };
+                // 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.
+                // Bail out.
+                if !assign_dominates {
+                    *set = Set1::Many;
+                }
+            }
+            PlaceContext::NonUse(_) => {}
         }
     }
 }

tests/mir-opt/copy-prop/non_dominate.f.CopyProp.diff

@@ -0,0 +1,29 @@
+- // MIR for `f` before CopyProp
++ // MIR for `f` after CopyProp
+  
+  fn f(_1: bool) -> bool {
+      let mut _0: bool;                    // return place in scope 0 at $DIR/non_dominate.rs:+0:18: +0:22
+      let mut _2: bool;                    // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+      let mut _3: bool;                    // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+  
+      bb0: {
+          goto -> bb1;                     // scope 0 at $DIR/non_dominate.rs:+4:11: +4:20
+      }
+  
+      bb1: {
+          _3 = _1;                         // scope 0 at $DIR/non_dominate.rs:+5:17: +5:22
+          switchInt(_3) -> [0: bb3, otherwise: bb2]; // scope 0 at $DIR/non_dominate.rs:+5:24: +5:58
+      }
+  
+      bb2: {
+          _2 = _3;                         // scope 0 at $DIR/non_dominate.rs:+8:17: +8:22
+          _1 = const false;                // scope 0 at $DIR/non_dominate.rs:+8:24: +8:33
+          goto -> bb1;                     // scope 0 at $DIR/non_dominate.rs:+8:35: +8:44
+      }
+  
+      bb3: {
+          _0 = _2;                         // scope 0 at $DIR/non_dominate.rs:+9:17: +9:24
+          return;                          // scope 0 at $DIR/non_dominate.rs:+9:26: +9:34
+      }
+  }
+  

tests/mir-opt/copy-prop/non_dominate.rs

@@ -0,0 +1,26 @@
+// unit-test: CopyProp
+
+#![feature(custom_mir, core_intrinsics)]
+#![allow(unused_assignments)]
+extern crate core;
+use core::intrinsics::mir::*;
+
+#[custom_mir(dialect = "analysis", phase = "post-cleanup")]
+fn f(c: bool) -> bool {
+    mir!(
+        let a: bool;
+        let b: bool;
+        { Goto(bb1) }
+        bb1 = { b = c; match b { false => bb3, _ => bb2 }}
+        // This assignment to `a` does not dominate the use in `bb3`.
+        // It should not be replaced by `b`.
+        bb2 = { a = b; c = false; Goto(bb1) }
+        bb3 = { RET = a; Return() }
+    )
+}
+
+fn main() {
+    assert_eq!(true, f(true));
+}
+
+// EMIT_MIR non_dominate.f.CopyProp.diff