GVN: Invalidate derefs when a loop exists

Author: dianqk

compiler/rustc_middle/src/mir/loop.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,8 @@ mod statement;
 mod syntax;
 mod terminator;
 
+#[path = "loop.rs"]
+pub mod loop_;
 pub mod traversal;
 pub mod visit;
 

compiler/rustc_mir_transform/src/gvn.rs

@@ -129,6 +129,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 = !loop_::maybe_loop_headers(body).is_empty();
 
         let arena = DroplessArena::default();
         let mut state =
@@ -141,6 +142,10 @@ 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.
+            if maybe_loop {
+                state.invalidate_derefs();
+            }
             let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
             state.visit_basic_block_data(bb, data);
         }
@@ -1924,7 +1929,9 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
         // Currently, only preserving derefs for trivial terminators like SwitchInt and Goto.
         let safe_to_preserve_derefs = matches!(
             terminator.kind,
-            TerminatorKind::SwitchInt { .. } | TerminatorKind::Goto { .. }
+            TerminatorKind::SwitchInt { .. }
+                | TerminatorKind::Goto { .. }
+                | TerminatorKind::Unreachable
         );
         if !safe_to_preserve_derefs {
             self.invalidate_derefs();

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),