Rollup merge of rust-lang#67130 - wesleywiser:const_prop_into_locals,…

… r=oli-obk

Const prop should finish propagation into user defined variables

Fixes rust-lang#66638

~~Temporarily rebased on top of rust-lang#67015 to get those fixes.~~

r? @oli-obk

src/librustc_mir/transform/const_prop.rs

@@ -262,7 +262,7 @@ struct ConstPropagator<'mir, 'tcx> {
     ecx: InterpCx<'mir, 'tcx, ConstPropMachine>,
     tcx: TyCtxt<'tcx>,
     source: MirSource<'tcx>,
-    can_const_prop: IndexVec<Local, bool>,
+    can_const_prop: IndexVec<Local, ConstPropMode>,
     param_env: ParamEnv<'tcx>,
     // FIXME(eddyb) avoid cloning these two fields more than once,
     // by accessing them through `ecx` instead.
@@ -708,17 +708,28 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
     }
 }
 
+/// The mode that `ConstProp` is allowed to run in for a given `Local`.
+#[derive(Clone, Copy, Debug, PartialEq)]
+enum ConstPropMode {
+    /// The `Local` can be propagated into and reads of this `Local` can also be propagated.
+    FullConstProp,
+    /// The `Local` can be propagated into but reads cannot be propagated.
+    OnlyPropagateInto,
+    /// No propagation is allowed at all.
+    NoPropagation,
+}
+
 struct CanConstProp {
-    can_const_prop: IndexVec<Local, bool>,
+    can_const_prop: IndexVec<Local, ConstPropMode>,
     // false at the beginning, once set, there are not allowed to be any more assignments
     found_assignment: IndexVec<Local, bool>,
 }
 
 impl CanConstProp {
     /// returns true if `local` can be propagated
-    fn check(body: ReadOnlyBodyAndCache<'_, '_>) -> IndexVec<Local, bool> {
+    fn check(body: ReadOnlyBodyAndCache<'_, '_>) -> IndexVec<Local, ConstPropMode> {
         let mut cpv = CanConstProp {
-            can_const_prop: IndexVec::from_elem(true, &body.local_decls),
+            can_const_prop: IndexVec::from_elem(ConstPropMode::FullConstProp, &body.local_decls),
             found_assignment: IndexVec::from_elem(false, &body.local_decls),
         };
         for (local, val) in cpv.can_const_prop.iter_enumerated_mut() {
@@ -728,10 +739,10 @@ impl CanConstProp {
             // FIXME(oli-obk): lint variables until they are used in a condition
             // FIXME(oli-obk): lint if return value is constant
             let local_kind = body.local_kind(local);
-            *val = local_kind == LocalKind::Temp || local_kind == LocalKind::ReturnPointer;
 
-            if !*val {
-                trace!("local {:?} can't be propagated because it's not a temporary", local);
+            if local_kind == LocalKind::Arg || local_kind == LocalKind::Var {
+                *val = ConstPropMode::OnlyPropagateInto;
+                trace!("local {:?} can't be const propagated because it's not a temporary", local);
             }
         }
         cpv.visit_body(body);
@@ -753,7 +764,7 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
             // only occur in independent execution paths
             MutatingUse(MutatingUseContext::Store) => if self.found_assignment[local] {
                 trace!("local {:?} can't be propagated because of multiple assignments", local);
-                self.can_const_prop[local] = false;
+                self.can_const_prop[local] = ConstPropMode::NoPropagation;
             } else {
                 self.found_assignment[local] = true
             },
@@ -766,7 +777,7 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
             NonUse(_) => {},
             _ => {
                 trace!("local {:?} can't be propagaged because it's used: {:?}", local, context);
-                self.can_const_prop[local] = false;
+                self.can_const_prop[local] = ConstPropMode::NoPropagation;
             },
         }
     }
@@ -800,10 +811,10 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
             if let Ok(place_layout) = self.tcx.layout_of(self.param_env.and(place_ty)) {
                 if let Some(local) = place.as_local() {
                     let source = statement.source_info;
+                    let can_const_prop = self.can_const_prop[local];
                     if let Some(()) = self.const_prop(rval, place_layout, source, place) {
-                        if self.can_const_prop[local] {
-                            trace!("propagated into {:?}", local);
-
+                        if can_const_prop == ConstPropMode::FullConstProp ||
+                           can_const_prop == ConstPropMode::OnlyPropagateInto {
                             if let Some(value) = self.get_const(local) {
                                 if self.should_const_prop(value) {
                                     trace!("replacing {:?} with {:?}", rval, value);
@@ -812,21 +823,26 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
                                         value,
                                         statement.source_info,
                                     );
+
+                                    if can_const_prop == ConstPropMode::FullConstProp {
+                                        trace!("propagated into {:?}", local);
+                                    }
                                 }
                             }
-                        } else {
-                            trace!("can't propagate into {:?}", local);
-                            if local != RETURN_PLACE {
-                                self.remove_const(local);
-                            }
+                        }
+                    }
+                    if self.can_const_prop[local] != ConstPropMode::FullConstProp {
+                        trace!("can't propagate into {:?}", local);
+                        if local != RETURN_PLACE {
+                            self.remove_const(local);
                         }
                     }
                 }
             }
         } else {
             match statement.kind {
                 StatementKind::StorageLive(local) |
-                StatementKind::StorageDead(local) if self.can_const_prop[local] => {
+                StatementKind::StorageDead(local) => {
                     let frame = self.ecx.frame_mut();
                     frame.locals[local].value =
                         if let StatementKind::StorageLive(_) = statement.kind {

src/test/mir-opt/const_prop/aggregate.rs

@@ -19,7 +19,7 @@ fn main() {
 //      ...
 //      _3 = (const 0i32, const 1i32, const 2i32);
 //      _2 = const 1i32;
-//      _1 = Add(move _2, const 0i32);
+//      _1 = const 1i32;
 //      ...
 //  }
 // END rustc.main.ConstProp.after.mir

src/test/mir-opt/const_prop/array_index.rs

@@ -26,7 +26,7 @@ fn main() {
 //      assert(const true, "index out of bounds: the len is move _4 but the index is _3") -> bb1;
 //  }
 //  bb1: {
-//      _1 = _2[_3];
+//      _1 = const 2u32;
 //      ...
 //      return;
 //  }

src/test/mir-opt/const_prop/optimizes_into_variable.rs

@@ -0,0 +1,149 @@
+// compile-flags: -C overflow-checks=on
+
+struct Point {
+    x: u32,
+    y: u32,
+}
+
+fn main() {
+    let x = 2 + 2;
+    let y = [0, 1, 2, 3, 4, 5][3];
+    let z = (Point { x: 12, y: 42}).y;
+}
+
+// END RUST SOURCE
+// START rustc.main.ConstProp.before.mir
+// let mut _0: ();
+// let _1: i32;
+// let mut _2: (i32, bool);
+// let mut _4: [i32; 6];
+// let _5: usize;
+// let mut _6: usize;
+// let mut _7: bool;
+// let mut _9: Point;
+// scope 1 {
+//   debug x => _1;
+//   let _3: i32;
+//   scope 2 {
+//     debug y => _3;
+//     let _8: u32;
+//     scope 3 {
+//       debug z => _8;
+//     }
+//   }
+// }
+// bb0: {
+//   StorageLive(_1);
+//   _2 = CheckedAdd(const 2i32, const 2i32);
+//   assert(!move (_2.1: bool), "attempt to add with overflow") -> bb1;
+// }
+// bb1: {
+//   _1 = move (_2.0: i32);
+//   StorageLive(_3);
+//   StorageLive(_4);
+//   _4 = [const 0i32, const 1i32, const 2i32, const 3i32, const 4i32, const 5i32];
+//   StorageLive(_5);
+//   _5 = const 3usize;
+//   _6 = const 6usize;
+//   _7 = Lt(_5, _6);
+//   assert(move _7, "index out of bounds: the len is move _6 but the index is _5") -> bb2;
+// }
+// bb2: {
+//   _3 = _4[_5];
+//   StorageDead(_5);
+//   StorageDead(_4);
+//   StorageLive(_8);
+//   StorageLive(_9);
+//   _9 = Point { x: const 12u32, y: const 42u32 };
+//   _8 = (_9.1: u32);
+//   StorageDead(_9);
+//   _0 = ();
+//   StorageDead(_8);
+//   StorageDead(_3);
+//   StorageDead(_1);
+//   return;
+// }
+// END rustc.main.ConstProp.before.mir
+// START rustc.main.ConstProp.after.mir
+// let mut _0: ();
+// let _1: i32;
+// let mut _2: (i32, bool);
+// let mut _4: [i32; 6];
+// let _5: usize;
+// let mut _6: usize;
+// let mut _7: bool;
+// let mut _9: Point;
+// scope 1 {
+//   debug x => _1;
+//   let _3: i32;
+//   scope 2 {
+//     debug y => _3;
+//     let _8: u32;
+//     scope 3 {
+//       debug z => _8;
+//     }
+//   }
+// }
+// bb0: {
+//   StorageLive(_1);
+//   _2 = (const 4i32, const false);
+//   assert(!const false, "attempt to add with overflow") -> bb1;
+// }
+// bb1: {
+//   _1 = const 4i32;
+//   StorageLive(_3);
+//   StorageLive(_4);
+//   _4 = [const 0i32, const 1i32, const 2i32, const 3i32, const 4i32, const 5i32];
+//   StorageLive(_5);
+//   _5 = const 3usize;
+//   _6 = const 6usize;
+//   _7 = const true;
+//   assert(const true, "index out of bounds: the len is move _6 but the index is _5") -> bb2;
+// }
+// bb2: {
+//   _3 = const 3i32;
+//   StorageDead(_5);
+//   StorageDead(_4);
+//   StorageLive(_8);
+//   StorageLive(_9);
+//   _9 = Point { x: const 12u32, y: const 42u32 };
+//   _8 = const 42u32;
+//   StorageDead(_9);
+//   _0 = ();
+//   StorageDead(_8);
+//   StorageDead(_3);
+//   StorageDead(_1);
+//   return;
+// }
+// END rustc.main.ConstProp.after.mir
+// START rustc.main.SimplifyLocals.after.mir
+// let mut _0: ();
+// let _1: i32;
+// let mut _3: [i32; 6];
+// scope 1 {
+//   debug x => _1;
+//   let _2: i32;
+//   scope 2 {
+//     debug y => _2;
+//     let _4: u32;
+//     scope 3 {
+//       debug z => _4;
+//     }
+//   }
+// }
+// bb0: {
+//   StorageLive(_1);
+//   _1 = const 4i32;
+//   StorageLive(_2);
+//   StorageLive(_3);
+//   _3 = [const 0i32, const 1i32, const 2i32, const 3i32, const 4i32, const 5i32];
+//   _2 = const 3i32;
+//   StorageDead(_3);
+//   StorageLive(_4);
+//   _4 = const 42u32;
+//   StorageDead(_4);
+//   StorageDead(_2);
+//   StorageDead(_1);
+//   return;
+// }
+// END rustc.main.SimplifyLocals.after.mir

src/test/mir-opt/const_prop/read_immutable_static.rs

@@ -25,7 +25,7 @@ fn main() {
 //      _2 = const 2u8;
 //      ...
 //      _4 = const 2u8;
-//      _1 = Add(move _2, move _4);
+//      _1 = const 4u8;
 //      ...
 //  }
 // END rustc.main.ConstProp.after.mir

src/test/mir-opt/const_prop/repeat.rs

@@ -30,7 +30,7 @@ fn main() {
 //  }
 //  bb1: {
 //      _2 = const 42u32;
-//      _1 = Add(move _2, const 0u32);
+//      _1 = const 42u32;
 //      ...
 //      return;
 //  }