Propagate PlaceElem::Index.

Author: cjgillot

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -6,7 +6,7 @@ use rustc_const_eval::const_eval::CheckAlignment;
 use rustc_const_eval::interpret::{ConstValue, ImmTy, Immediate, InterpCx, Scalar};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def::DefKind;
-use rustc_middle::mir::visit::{MutVisitor, Visitor};
+use rustc_middle::mir::visit::{MutVisitor, NonMutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -579,11 +579,29 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx> {
             Operand::Copy(place) | Operand::Move(place) => {
                 if let Some(value) = self.before_effect.get(&(location, *place)) {
                     *operand = self.make_operand(value.clone());
+                } else if !place.projection.is_empty() {
+                    self.super_operand(operand, location)
                 }
             }
             _ => (),
         }
     }
+
+    fn process_projection_elem(
+        &mut self,
+        elem: PlaceElem<'tcx>,
+        location: Location,
+    ) -> Option<PlaceElem<'tcx>> {
+        if let PlaceElem::Index(local) = elem
+            && let Some(value) = self.before_effect.get(&(location, local.into()))
+            && let Ok(offset) = value.0.to_target_usize(&self.tcx)
+            && let Some(min_length) = offset.checked_add(1)
+        {
+            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
+        } else {
+            None
+        }
+    }
 }
 
 pub(crate) struct OperandCollector<'tcx, 'map, 'a> {
@@ -594,17 +612,21 @@ pub(crate) struct OperandCollector<'tcx, 'map, 'a> {
 
 impl<'tcx, 'map, 'a> Visitor<'tcx> for OperandCollector<'tcx, 'map, 'a> {
     fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
-        match operand {
-            Operand::Copy(place) | Operand::Move(place) => {
-                match self.state.get(place.as_ref(), self.map) {
-                    FlatSet::Top => (),
-                    FlatSet::Elem(value) => {
-                        self.visitor.before_effect.insert((location, *place), value);
-                    }
-                    FlatSet::Bottom => (),
-                }
+        if let Some(place) = operand.place() {
+            if let FlatSet::Elem(value) = self.state.get(place.as_ref(), self.map) {
+                self.visitor.before_effect.insert((location, place), value);
+            } else if !place.projection.is_empty() {
+                // Try to propagate into `Index` projections.
+                self.super_operand(operand, location)
             }
-            _ => (),
+        }
+    }
+
+    fn visit_local(&mut self, local: Local, ctxt: PlaceContext, location: Location) {
+        if let PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy | NonMutatingUseContext::Move) = ctxt
+            && let FlatSet::Elem(value) = self.state.get(local.into(), self.map)
+        {
+            self.visitor.before_effect.insert((location, local.into()), value);
         }
     }
 }

tests/mir-opt/const_prop/array_index.main.ConstProp.32bit.panic-abort.diff

@@ -27,7 +27,8 @@
       }
 
       bb1: {
-          _1 = _2[_3];
+-         _1 = _2[_3];
++         _1 = _2[2 of 3];
           StorageDead(_3);
           StorageDead(_2);
           _0 = const ();

tests/mir-opt/const_prop/array_index.main.ConstProp.32bit.panic-unwind.diff

@@ -27,7 +27,8 @@
       }
 
       bb1: {
-          _1 = _2[_3];
+-         _1 = _2[_3];
++         _1 = _2[2 of 3];
           StorageDead(_3);
           StorageDead(_2);
           _0 = const ();

tests/mir-opt/const_prop/array_index.main.ConstProp.64bit.panic-abort.diff

@@ -27,7 +27,8 @@
       }
 
       bb1: {
-          _1 = _2[_3];
+-         _1 = _2[_3];
++         _1 = _2[2 of 3];
           StorageDead(_3);
           StorageDead(_2);
           _0 = const ();

tests/mir-opt/const_prop/array_index.main.ConstProp.64bit.panic-unwind.diff

@@ -27,7 +27,8 @@
       }
 
       bb1: {
-          _1 = _2[_3];
+-         _1 = _2[_3];
++         _1 = _2[2 of 3];
           StorageDead(_3);
           StorageDead(_2);
           _0 = const ();

tests/mir-opt/const_prop/bad_op_unsafe_oob_for_slices.main.ConstProp.32bit.panic-abort.diff

@@ -43,7 +43,8 @@
       }
 
       bb1: {
-          _5 = (*_1)[_6];
+-         _5 = (*_1)[_6];
++         _5 = (*_1)[3 of 4];
           StorageDead(_6);
           _0 = const ();
           StorageDead(_5);

tests/mir-opt/const_prop/bad_op_unsafe_oob_for_slices.main.ConstProp.32bit.panic-unwind.diff

@@ -43,7 +43,8 @@
       }
 
       bb1: {
-          _5 = (*_1)[_6];
+-         _5 = (*_1)[_6];
++         _5 = (*_1)[3 of 4];
           StorageDead(_6);
           _0 = const ();
           StorageDead(_5);

tests/mir-opt/const_prop/bad_op_unsafe_oob_for_slices.main.ConstProp.64bit.panic-abort.diff

@@ -43,7 +43,8 @@
       }
 
       bb1: {
-          _5 = (*_1)[_6];
+-         _5 = (*_1)[_6];
++         _5 = (*_1)[3 of 4];
           StorageDead(_6);
           _0 = const ();
           StorageDead(_5);

tests/mir-opt/const_prop/bad_op_unsafe_oob_for_slices.main.ConstProp.64bit.panic-unwind.diff

@@ -43,7 +43,8 @@
       }
 
       bb1: {
-          _5 = (*_1)[_6];
+-         _5 = (*_1)[_6];
++         _5 = (*_1)[3 of 4];
           StorageDead(_6);
           _0 = const ();
           StorageDead(_5);

tests/mir-opt/const_prop/large_array_index.main.ConstProp.32bit.panic-abort.diff

@@ -27,7 +27,8 @@
       }
 
       bb1: {
-          _1 = _2[_3];
+-         _1 = _2[_3];
++         _1 = _2[2 of 3];
           StorageDead(_3);
           StorageDead(_2);
           _0 = const ();