GVN: Invalidate derefs at loop headers

(cherry picked from commit 2048b9c)

Author: dianqk

compiler/rustc_middle/src/mir/loops.rs

@@ -0,0 +1,29 @@
+use rustc_index::bit_set::DenseBitSet;
+
+use super::*;
+
+/// Compute the set of loop headers in the given body. A loop header is usually defined as a block
+/// which dominates one of its predecessors. This definition is only correct for reducible CFGs.
+/// However, computing dominators is expensive, so we approximate according to the post-order
+/// traversal order. A loop header for us is a block which is visited after its predecessor in
+/// post-order. This is ok as we mostly need a heuristic.
+pub fn maybe_loop_headers(body: &Body<'_>) -> DenseBitSet<BasicBlock> {
+    let mut maybe_loop_headers = DenseBitSet::new_empty(body.basic_blocks.len());
+    let mut visited = DenseBitSet::new_empty(body.basic_blocks.len());
+    for (bb, bbdata) in traversal::postorder(body) {
+        // Post-order means we visit successors before the block for acyclic CFGs.
+        // If the successor is not visited yet, consider it a loop header.
+        for succ in bbdata.terminator().successors() {
+            if !visited.contains(succ) {
+                maybe_loop_headers.insert(succ);
+            }
+        }
+
+        // Only mark `bb` as visited after we checked the successors, in case we have a self-loop.
+        //     bb1: goto -> bb1;
+        let _new = visited.insert(bb);
+        debug_assert!(_new);
+    }
+
+    maybe_loop_headers
+}

compiler/rustc_middle/src/mir/mod.rs

@@ -51,6 +51,7 @@ mod statement;
 mod syntax;
 mod terminator;
 
+pub mod loops;
 pub mod traversal;
 pub mod visit;
 

compiler/rustc_mir_transform/src/gvn.rs

@@ -126,6 +126,7 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
         let ssa = SsaLocals::new(tcx, body, typing_env);
         // Clone dominators because we need them while mutating the body.
         let dominators = body.basic_blocks.dominators().clone();
+        let maybe_loop_headers = loops::maybe_loop_headers(body);
 
         let mut state = VnState::new(tcx, body, typing_env, &ssa, dominators, &body.local_decls);
 
@@ -136,6 +137,11 @@ impl<'tcx> crate::MirPass<'tcx> for GVN {
 
         let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec();
         for bb in reverse_postorder {
+            // N.B. With loops, reverse postorder cannot produce a valid topological order.
+            // A statement or terminator from inside the loop, that is not processed yet, may have performed an indirect write.
+            if maybe_loop_headers.contains(bb) {
+                state.invalidate_derefs();
+            }
             let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
             state.visit_basic_block_data(bb, data);
         }

compiler/rustc_mir_transform/src/jump_threading.rs

@@ -84,7 +84,7 @@ impl<'tcx> crate::MirPass<'tcx> for JumpThreading {
             body,
             arena,
             map: Map::new(tcx, body, Some(MAX_PLACES)),
-            loop_headers: loop_headers(body),
+            loop_headers: loops::maybe_loop_headers(body),
             opportunities: Vec::new(),
         };
 
@@ -832,21 +832,3 @@ enum Update {
     Incr,
     Decr,
 }
-
-/// Compute the set of loop headers in the given body. We define a loop header as a block which has
-/// at least a predecessor which it dominates. This definition is only correct for reducible CFGs.
-/// But if the CFG is already irreducible, there is no point in trying much harder.
-/// is already irreducible.
-fn loop_headers(body: &Body<'_>) -> DenseBitSet<BasicBlock> {
-    let mut loop_headers = DenseBitSet::new_empty(body.basic_blocks.len());
-    let dominators = body.basic_blocks.dominators();
-    // Only visit reachable blocks.
-    for (bb, bbdata) in traversal::preorder(body) {
-        for succ in bbdata.terminator().successors() {
-            if dominators.dominates(succ, bb) {
-                loop_headers.insert(succ);
-            }
-        }
-    }
-    loop_headers
-}

tests/mir-opt/gvn_loop.loop_deref_mut.GVN.diff

@@ -71,7 +71,7 @@
           StorageLive(_12);
           _12 = copy _7;
 -         _11 = Value::V0(move _12);
-+         _11 = copy (*_3);
++         _11 = Value::V0(copy _7);
           StorageDead(_12);
           StorageLive(_13);
           StorageLive(_14);

tests/mir-opt/gvn_loop.rs

@@ -1,4 +1,3 @@
-// skip-filecheck
 //@ test-mir-pass: GVN
 
 #![crate_type = "lib"]
@@ -9,8 +8,17 @@ pub enum Value {
     V1,
 }
 
+// Check that we do not use the dereferenced value of `val_alias` when returning.
+
 // EMIT_MIR gvn_loop.loop_deref_mut.GVN.diff
 fn loop_deref_mut(val: &mut Value) -> Value {
+    // CHECK-LABEL: fn loop_deref_mut(
+    // CHECK: [[VAL_REF:_.*]] = get::<Value>(
+    // CHECK: [[V:_.*]] = copy (((*[[VAL_REF]]) as V0).0: i32);
+    // CEHCK-NOT: copy (*[[VAL_REF]]);
+    // CHECK: [[RET:_*]] = Value::V0(copy [[V]]);
+    // CEHCK-NOT: copy (*[[VAL_REF]]);
+    // CHECK: _0 = move [[RET]]
     let val_alias: &Value = get(val);
     let mut stop = false;
     let Value::V0(v) = *val_alias else { unsafe { core::intrinsics::unreachable() } };

tests/ui/mir/gvn-loop-miscompile.rs

@@ -1,5 +1,5 @@
 //@ compile-flags: -O
-//@ run-fail
+//@ run-pass
 
 pub enum Value {
     V0(i32),