Rollup merge of #141252 - dianqk:gvn-repeat-index, r=saethlin

gvn: bail out unavoidable non-ssa locals in repeat

Fixes #141251.

We cannot transform `*elem` to `array[idx1]` in the following code, as `idx1` has already been modified.

```rust
    mir! {
        let array;
        let elem;
        {
            array = [*val; 5];
            elem = &array[idx1];
            idx1 = idx2;
            RET = *elem;
            Return()
        }
    }
```

Perhaps I could transform it to `array[0]`, but I prefer the conservative approach.

r? mir-opt

4 files changed, 91 insertions, 1 deletions

diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
index 209e818e9e3..a91d46ec406 100644
--- a/compiler/rustc_mir_transform/src/gvn.rs
+++ b/compiler/rustc_mir_transform/src/gvn.rs
@@ -638,6 +638,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         place: PlaceRef<'tcx>,
         value: VnIndex,
         proj: PlaceElem<'tcx>,
+        from_non_ssa_index: &mut bool,
     ) -> Option<VnIndex> {
         let proj = match proj {
             ProjectionElem::Deref => {
@@ -682,6 +683,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             }
             ProjectionElem::Index(idx) => {
                 if let Value::Repeat(inner, _) = self.get(value) {
+                    *from_non_ssa_index |= self.locals[idx].is_none();
                     return Some(*inner);
                 }
                 let idx = self.locals[idx]?;
@@ -774,6 +776,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
 
         // Invariant: `value` holds the value up-to the `index`th projection excluded.
         let mut value = self.locals[place.local]?;
+        let mut from_non_ssa_index = false;
         for (index, proj) in place.projection.iter().enumerate() {
             if let Value::Projection(pointer, ProjectionElem::Deref) = *self.get(value)
                 && let Value::Address { place: mut pointee, kind, .. } = *self.get(pointer)
@@ -791,7 +794,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             }
 
             let base = PlaceRef { local: place.local, projection: &place.projection[..index] };
-            value = self.project(base, value, proj)?;
+            value = self.project(base, value, proj, &mut from_non_ssa_index)?;
         }
 
         if let Value::Projection(pointer, ProjectionElem::Deref) = *self.get(value)
@@ -804,6 +807,9 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
         }
         if let Some(new_local) = self.try_as_local(value, location) {
             place_ref = PlaceRef { local: new_local, projection: &[] };
+        } else if from_non_ssa_index {
+            // If access to non-SSA locals is unavoidable, bail out.
+            return None;
         }
 
         if place_ref.local != place.local || place_ref.projection.len() < place.projection.len() {
diff --git a/tests/mir-opt/gvn_repeat.repeat_local.GVN.diff b/tests/mir-opt/gvn_repeat.repeat_local.GVN.diff
new file mode 100644
index 00000000000..fd047825281
--- /dev/null
+++ b/tests/mir-opt/gvn_repeat.repeat_local.GVN.diff
@@ -0,0 +1,18 @@
+- // MIR for `repeat_local` before GVN
++ // MIR for `repeat_local` after GVN
+  
+  fn repeat_local(_1: usize, _2: usize, _3: i32) -> i32 {
+      let mut _0: i32;
+      let mut _4: [i32; 5];
+      let mut _5: &i32;
+  
+      bb0: {
+          _4 = [copy _3; 5];
+          _5 = &_4[_1];
+          _1 = copy _2;
+-         _0 = copy (*_5);
++         _0 = copy _3;
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/gvn_repeat.repeat_place.GVN.diff b/tests/mir-opt/gvn_repeat.repeat_place.GVN.diff
new file mode 100644
index 00000000000..e490925bc11
--- /dev/null
+++ b/tests/mir-opt/gvn_repeat.repeat_place.GVN.diff
@@ -0,0 +1,17 @@
+- // MIR for `repeat_place` before GVN
++ // MIR for `repeat_place` after GVN
+  
+  fn repeat_place(_1: usize, _2: usize, _3: &i32) -> i32 {
+      let mut _0: i32;
+      let mut _4: [i32; 5];
+      let mut _5: &i32;
+  
+      bb0: {
+          _4 = [copy (*_3); 5];
+          _5 = &_4[_1];
+          _1 = copy _2;
+          _0 = copy (*_5);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/gvn_repeat.rs b/tests/mir-opt/gvn_repeat.rs
new file mode 100644
index 00000000000..bbbb2a7ccba
--- /dev/null
+++ b/tests/mir-opt/gvn_repeat.rs
@@ -0,0 +1,49 @@
+//@ test-mir-pass: GVN
+
+#![feature(custom_mir, core_intrinsics)]
+
+// Check that we do not introduce out-of-bounds access.
+
+use std::intrinsics::mir::*;
+
+// EMIT_MIR gvn_repeat.repeat_place.GVN.diff
+#[custom_mir(dialect = "runtime")]
+pub fn repeat_place(mut idx1: usize, idx2: usize, val: &i32) -> i32 {
+    // CHECK-LABEL: fn repeat_place(
+    // CHECK: let mut [[ELEM:.*]]: &i32;
+    // CHECK: _0 = copy (*[[ELEM]])
+    mir! {
+        let array;
+        let elem;
+        {
+            array = [*val; 5];
+            elem = &array[idx1];
+            idx1 = idx2;
+            RET = *elem;
+            Return()
+        }
+    }
+}
+
+// EMIT_MIR gvn_repeat.repeat_local.GVN.diff
+#[custom_mir(dialect = "runtime")]
+pub fn repeat_local(mut idx1: usize, idx2: usize, val: i32) -> i32 {
+    // CHECK-LABEL: fn repeat_local(
+    // CHECK: _0 = copy _3
+    mir! {
+        let array;
+        let elem;
+        {
+            array = [val; 5];
+            elem = &array[idx1];
+            idx1 = idx2;
+            RET = *elem;
+            Return()
+        }
+    }
+}
+
+fn main() {
+    assert_eq!(repeat_place(0, 5, &0), 0);
+    assert_eq!(repeat_local(0, 5, 0), 0);
+}