Propagate PlaceElem::Index.

cjgillot · cjgillot · commit f96c6e04cb93 · 2023-09-06T16:09:31.000Z
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/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, ScalarInt, Ty, TyCtxt};
@@ -545,11 +545,29 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
                 if let Some(value) = self.before_effect.get(&(location, *place)) {
                     let ty = place.ty(self.local_decls, self.tcx).ty;
                     *operand = self.make_operand(*value, ty);
+                } else if !place.projection.is_empty() {
+                    self.super_operand(operand, location)
                 }
             }
             Operand::Constant(_) => {}
         }
     }
+
+    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.try_to_target_usize(self.tcx)
+            && let Some(min_length) = offset.checked_add(1)
+        {
+            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
+        } else {
+            None
+        }
+    }
 }
 
 struct OperandCollector<'tcx, 'map, 'locals, 'a> {
@@ -560,17 +578,21 @@ struct OperandCollector<'tcx, 'map, 'locals, 'a> {
 
 impl<'tcx> Visitor<'tcx> for OperandCollector<'tcx, '_, '_, '_> {
     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);
         }
     }
 }
diff --git a/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.32bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.32bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.64bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.64bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.32bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.32bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.64bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.64bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/large_array_index.main.DataflowConstProp.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 ();
diff --git a/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.32bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.32bit.panic-abort.diff
@@ -29,7 +29,8 @@
       }
   
       bb1: {
-          _2 = _3[_4];
+-         _2 = _3[_4];
++         _2 = _3[2 of 3];
           _1 = Add(move _2, const 0_u32);
           StorageDead(_2);
           StorageDead(_4);
diff --git a/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.32bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.32bit.panic-unwind.diff
@@ -29,7 +29,8 @@
       }
   
       bb1: {
-          _2 = _3[_4];
+-         _2 = _3[_4];
++         _2 = _3[2 of 3];
           _1 = Add(move _2, const 0_u32);
           StorageDead(_2);
           StorageDead(_4);
diff --git a/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.64bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.64bit.panic-abort.diff
@@ -29,7 +29,8 @@
       }
   
       bb1: {
-          _2 = _3[_4];
+-         _2 = _3[_4];
++         _2 = _3[2 of 3];
           _1 = Add(move _2, const 0_u32);
           StorageDead(_2);
           StorageDead(_4);
diff --git a/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.64bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/repeat.main.DataflowConstProp.64bit.panic-unwind.diff
@@ -29,7 +29,8 @@
       }
   
       bb1: {
-          _2 = _3[_4];
+-         _2 = _3[_4];
++         _2 = _3[2 of 3];
           _1 = Add(move _2, const 0_u32);
           StorageDead(_2);
           StorageDead(_4);
diff --git a/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.32bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.32bit.panic-abort.diff
@@ -34,7 +34,8 @@
       }
   
       bb1: {
-          _1 = (*_2)[_6];
+-         _1 = (*_2)[_6];
++         _1 = (*_2)[1 of 2];
           StorageDead(_6);
           StorageDead(_4);
           StorageDead(_2);
diff --git a/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.32bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.32bit.panic-unwind.diff
@@ -34,7 +34,8 @@
       }
   
       bb1: {
-          _1 = (*_2)[_6];
+-         _1 = (*_2)[_6];
++         _1 = (*_2)[1 of 2];
           StorageDead(_6);
           StorageDead(_4);
           StorageDead(_2);
diff --git a/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.64bit.panic-abort.diff b/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.64bit.panic-abort.diff
@@ -34,7 +34,8 @@
       }
   
       bb1: {
-          _1 = (*_2)[_6];
+-         _1 = (*_2)[_6];
++         _1 = (*_2)[1 of 2];
           StorageDead(_6);
           StorageDead(_4);
           StorageDead(_2);
diff --git a/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.64bit.panic-unwind.diff b/tests/mir-opt/dataflow-const-prop/slice_len.main.DataflowConstProp.64bit.panic-unwind.diff
@@ -34,7 +34,8 @@
       }
   
       bb1: {
-          _1 = (*_2)[_6];
+-         _1 = (*_2)[_6];
++         _1 = (*_2)[1 of 2];
           StorageDead(_6);
           StorageDead(_4);
           StorageDead(_2);

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,8 @@`
`27`	`27`	`}`
`28`	`28`
`29`	`29`	`bb1: {`
`30`		`- _1 = _2[_3];`
	`30`	`+- _1 = _2[_3];`
	`31`	`++ _1 = _2[2 of 3];`
`31`	`32`	`StorageDead(_3);`
`32`	`33`	`StorageDead(_2);`
`33`	`34`	`_0 = const ();`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,8 @@`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`bb1: {`
`32`		`- _2 = _3[_4];`
	`32`	`+- _2 = _3[_4];`
	`33`	`++ _2 = _3[2 of 3];`
`33`	`34`	`_1 = Add(move _2, const 0_u32);`
`34`	`35`	`StorageDead(_2);`
`35`	`36`	`StorageDead(_4);`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,8 @@`
`34`	`34`	`}`
`35`	`35`
`36`	`36`	`bb1: {`
`37`		`- _1 = (*_2)[_6];`
	`37`	`+- _1 = (*_2)[_6];`
	`38`	`++ _1 = (*_2)[1 of 2];`
`38`	`39`	`StorageDead(_6);`
`39`	`40`	`StorageDead(_4);`
`40`	`41`	`StorageDead(_2);`