GVN: Invalidate derefs at loop headers

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

@@ -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_headers = loops::maybe_loop_headers(body);
 
         let arena = DroplessArena::default();
         let mut state =
@@ -141,6 +142,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)),
-            maybe_loop_headers: maybe_loop_headers(body),
+            maybe_loop_headers: loops::maybe_loop_headers(body),
             opportunities: Vec::new(),
         };
 
@@ -830,29 +830,3 @@ enum Update {
     Incr,
     Decr,
 }
-
-/// 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.
-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
-}

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

@@ -0,0 +1,115 @@
+- // MIR for `loop_deref_mut` before GVN
++ // MIR for `loop_deref_mut` after GVN
+
+  fn loop_deref_mut(_1: &mut Value) -> Value {
+      debug val => _1;
+      let mut _0: Value;
+      let _2: &Value;
+      let _3: &Value;
+      let mut _4: &Value;
+      let mut _6: !;
+      let mut _8: isize;
+      let mut _9: !;
+      let mut _10: ();
+      let mut _12: i32;
+      let _13: ();
+      let mut _14: bool;
+      let mut _15: !;
+      let mut _16: Value;
+      scope 1 {
+          debug val_alias => _2;
+          let mut _5: bool;
+          scope 2 {
+              debug stop => _5;
+              let _7: i32;
+              scope 3 {
+                  debug v => _7;
+                  let _11: Value;
+                  scope 4 {
+                      debug v => _11;
+                  }
+              }
+          }
+      }
+
+      bb0: {
+          StorageLive(_2);
+-         StorageLive(_3);
++         nop;
+          StorageLive(_4);
+          _4 = &(*_1);
+          _3 = get::<Value>(move _4) -> [return: bb1, unwind unreachable];
+      }
+
+      bb1: {
+          _2 = &(*_3);
+          StorageDead(_4);
+-         StorageDead(_3);
++         nop;
+          StorageLive(_5);
+          _5 = const false;
+-         _8 = discriminant((*_2));
++         _8 = discriminant((*_3));
+          switchInt(move _8) -> [0: bb3, otherwise: bb2];
+      }
+
+      bb2: {
+          unreachable;
+      }
+
+      bb3: {
+-         StorageLive(_7);
+-         _7 = copy (((*_2) as V0).0: i32);
++         nop;
++         _7 = copy (((*_3) as V0).0: i32);
+          StorageLive(_9);
+          goto -> bb4;
+      }
+
+      bb4: {
+          StorageLive(_11);
+          StorageLive(_12);
+          _12 = copy _7;
+-         _11 = Value::V0(move _12);
++         _11 = Value::V0(copy _7);
+          StorageDead(_12);
+          StorageLive(_13);
+          StorageLive(_14);
+          _14 = copy _5;
+          switchInt(move _14) -> [0: bb6, otherwise: bb5];
+      }
+
+      bb5: {
+          _0 = move _11;
+          StorageDead(_14);
+          StorageDead(_13);
+          StorageDead(_11);
+          StorageDead(_9);
+-         StorageDead(_7);
++         nop;
+          StorageDead(_5);
+          StorageDead(_2);
+          return;
+      }
+
+      bb6: {
+          _13 = const ();
+          StorageDead(_14);
+          StorageDead(_13);
+          _5 = const true;
+          StorageLive(_16);
+-         _16 = Value::V1;
+-         (*_1) = move _16;
++         _16 = const Value::V1;
++         (*_1) = const Value::V1;
+          StorageDead(_16);
+          _10 = const ();
+          StorageDead(_11);
+          goto -> bb4;
+      }
++ }
++ 
++ ALLOC0 (size: 8, align: 4) {
++     01 00 00 00 __ __ __ __                         │ ....░░░░
+  }
+

tests/mir-opt/gvn_loop.rs

@@ -0,0 +1,39 @@
+//@ test-mir-pass: GVN
+
+#![crate_type = "lib"]
+#![feature(core_intrinsics, rustc_attrs)]
+
+pub enum Value {
+    V0(i32),
+    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() } };
+    loop {
+        let v = Value::V0(v);
+        if stop {
+            return v;
+        }
+        stop = true;
+        *val = Value::V1;
+    }
+}
+
+#[inline(never)]
+#[rustc_nounwind]
+fn get<T>(v: &T) -> &T {
+    v
+}

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

@@ -0,0 +1,34 @@
+//@ compile-flags: -O
+//@ run-pass
+
+pub enum Value {
+    V0(i32),
+    V1,
+}
+
+fn set_discriminant(val: &mut Value) -> Value {
+    let val_alias: &Value = get(val);
+    let mut stop = false;
+    let Value::V0(v) = *val_alias else {
+        unreachable!();
+    };
+    loop {
+        let v = Value::V0(v);
+        if stop {
+            return v;
+        }
+        stop = true;
+        *val = Value::V1;
+    }
+}
+
+fn main() {
+    let mut v = Value::V0(1);
+    let v = set_discriminant(&mut v);
+    assert!(matches!(v, Value::V0(1)));
+}
+
+#[inline(never)]
+fn get<T>(v: &T) -> &T {
+    v
+}