Auto merge of #71953 - oli-obk:const_prop_deaggregates, r=wesleywiser

Const prop aggregates even if partially or fully modified

r? @wesleywiser

cc @rust-lang/wg-mir-opt I'm moderately scared of this change, but I'm confident in having reviewed all the cases.

src/librustc_mir/transform/const_prop.rs

@@ -349,8 +349,8 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }
     }
 
-    fn get_const(&self, local: Local) -> Option<OpTy<'tcx>> {
-        let op = self.ecx.access_local(self.ecx.frame(), local, None).ok();
+    fn get_const(&self, place: Place<'tcx>) -> Option<OpTy<'tcx>> {
+        let op = self.ecx.eval_place_to_op(place, None).ok();
 
         // Try to read the local as an immediate so that if it is representable as a scalar, we can
         // handle it as such, but otherwise, just return the value as is.
@@ -772,13 +772,25 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
     fn visit_local(&mut self, &local: &Local, context: PlaceContext, _: Location) {
         use rustc_middle::mir::visit::PlaceContext::*;
         match context {
-            // Constants must have at most one write
-            // FIXME(oli-obk): we could be more powerful here, if the multiple writes
-            // only occur in independent execution paths
-            MutatingUse(MutatingUseContext::Store) => {
+            // Projections are fine, because `&mut foo.x` will be caught by
+            // `MutatingUseContext::Borrow` elsewhere.
+            MutatingUse(MutatingUseContext::Projection)
+            | MutatingUse(MutatingUseContext::Store) => {
                 if !self.found_assignment.insert(local) {
-                    trace!("local {:?} can't be propagated because of multiple assignments", local);
-                    self.can_const_prop[local] = ConstPropMode::NoPropagation;
+                    match &mut self.can_const_prop[local] {
+                        // If the local can only get propagated in its own block, then we don't have
+                        // to worry about multiple assignments, as we'll nuke the const state at the
+                        // end of the block anyway, and inside the block we overwrite previous
+                        // states as applicable.
+                        ConstPropMode::OnlyInsideOwnBlock => {}
+                        other => {
+                            trace!(
+                                "local {:?} can't be propagated because of multiple assignments",
+                                local,
+                            );
+                            *other = ConstPropMode::NoPropagation;
+                        }
+                    }
                 }
             }
             // Reading constants is allowed an arbitrary number of times
@@ -787,12 +799,6 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
             | NonMutatingUse(NonMutatingUseContext::Inspect)
             | NonMutatingUse(NonMutatingUseContext::Projection)
             | NonUse(_) => {}
-            // FIXME(felix91gr): explain the reasoning behind this
-            MutatingUse(MutatingUseContext::Projection) => {
-                if self.local_kinds[local] != LocalKind::Temp {
-                    self.can_const_prop[local] = ConstPropMode::NoPropagation;
-                }
-            }
             _ => {
                 trace!("local {:?} can't be propagaged because it's used: {:?}", local, context);
                 self.can_const_prop[local] = ConstPropMode::NoPropagation;
@@ -826,40 +832,50 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
         if let StatementKind::Assign(box (place, ref mut rval)) = statement.kind {
             let place_ty: Ty<'tcx> = place.ty(&self.local_decls, self.tcx).ty;
             if let Ok(place_layout) = self.tcx.layout_of(self.param_env.and(place_ty)) {
-                if let Some(local) = place.as_local() {
-                    let can_const_prop = self.can_const_prop[local];
-                    if let Some(()) = self.const_prop(rval, place_layout, source_info, place) {
-                        if can_const_prop != ConstPropMode::NoPropagation {
-                            // This will return None for Locals that are from other blocks,
-                            // so it should be okay to propagate from here on down.
-                            if let Some(value) = self.get_const(local) {
-                                if self.should_const_prop(value) {
-                                    trace!("replacing {:?} with {:?}", rval, value);
-                                    self.replace_with_const(rval, value, statement.source_info);
-                                    if can_const_prop == ConstPropMode::FullConstProp
-                                        || can_const_prop == ConstPropMode::OnlyInsideOwnBlock
-                                    {
-                                        trace!("propagated into {:?}", local);
-                                    }
-                                }
-                                if can_const_prop == ConstPropMode::OnlyInsideOwnBlock {
-                                    trace!(
-                                        "found local restricted to its block. Will remove it from const-prop after block is finished. Local: {:?}",
-                                        local
-                                    );
-                                    self.locals_of_current_block.insert(local);
+                let can_const_prop = self.can_const_prop[place.local];
+                if let Some(()) = self.const_prop(rval, place_layout, source_info, place) {
+                    if can_const_prop != ConstPropMode::NoPropagation {
+                        // This will return None for variables that are from other blocks,
+                        // so it should be okay to propagate from here on down.
+                        if let Some(value) = self.get_const(place) {
+                            if self.should_const_prop(value) {
+                                trace!("replacing {:?} with {:?}", rval, value);
+                                self.replace_with_const(rval, value, statement.source_info);
+                                if can_const_prop == ConstPropMode::FullConstProp
+                                    || can_const_prop == ConstPropMode::OnlyInsideOwnBlock
+                                {
+                                    trace!("propagated into {:?}", place);
                                 }
                             }
+                            if can_const_prop == ConstPropMode::OnlyInsideOwnBlock {
+                                trace!(
+                                    "found local restricted to its block. Will remove it from const-prop after block is finished. Local: {:?}",
+                                    place.local
+                                );
+                                self.locals_of_current_block.insert(place.local);
+                            }
                         }
                     }
-                    if self.can_const_prop[local] == ConstPropMode::OnlyPropagateInto
-                        || self.can_const_prop[local] == ConstPropMode::NoPropagation
+                    if can_const_prop == ConstPropMode::OnlyPropagateInto
+                        || can_const_prop == ConstPropMode::NoPropagation
                     {
-                        trace!("can't propagate into {:?}", local);
-                        if local != RETURN_PLACE {
-                            Self::remove_const(&mut self.ecx, local);
+                        trace!("can't propagate into {:?}", place);
+                        if place.local != RETURN_PLACE {
+                            Self::remove_const(&mut self.ecx, place.local);
                         }
                     }
+                } else {
+                    // Const prop failed, so erase the destination, ensuring that whatever happens
+                    // from here on, does not know about the previous value.
+                    // This is important in case we have
+                    // ```rust
+                    // let mut x = 42;
+                    // x = SOME_MUTABLE_STATIC;
+                    // // x must now be undefined
+                    // ```
+                    // FIXME: we overzealously erase the entire local, because that's easier to
+                    // implement.
+                    Self::remove_const(&mut self.ecx, place.local);
                 }
             }
         } else {
@@ -993,7 +1009,7 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
                           arguments are of the variant `Operand::Copy`. This allows us to
                           simplify our handling of `Operands` in this case.
                         */
-                        if let Some(l) = opr.place().and_then(|p| p.as_local()) {
+                        if let Some(l) = opr.place() {
                             if let Some(value) = self.get_const(l) {
                                 if self.should_const_prop(value) {
                                     // FIXME(felix91gr): this code only handles `Scalar` cases.

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

@@ -0,0 +1,8 @@
+// compile-flags: -O
+
+// EMIT_MIR rustc.main.ConstProp.diff
+fn main() {
+    let mut x = 42;
+    x = 99;
+    let y = x;
+}

src/test/mir-opt/const_prop/mutable_variable/rustc.main.ConstProp.diff

@@ -0,0 +1,53 @@
+- // MIR for `main` before ConstProp
++ // MIR for `main` after ConstProp
+
+  fn main() -> () {
+      let mut _0: ();                      // return place in scope 0 at $DIR/mutable_variable.rs:4:11: 4:11
+      let mut _1: i32;                     // in scope 0 at $DIR/mutable_variable.rs:5:9: 5:14
+      scope 1 {
+          debug x => _1;                   // in scope 1 at $DIR/mutable_variable.rs:5:9: 5:14
+          let _2: i32;                     // in scope 1 at $DIR/mutable_variable.rs:7:9: 7:10
+          scope 2 {
+              debug y => _2;               // in scope 2 at $DIR/mutable_variable.rs:7:9: 7:10
+          }
+      }
+
+      bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/mutable_variable.rs:5:9: 5:14
+          _1 = const 42i32;                // scope 0 at $DIR/mutable_variable.rs:5:17: 5:19
+                                           // ty::Const
+                                           // + ty: i32
+                                           // + val: Value(Scalar(0x0000002a))
+                                           // mir::Constant
+                                           // + span: $DIR/mutable_variable.rs:5:17: 5:19
+                                           // + literal: Const { ty: i32, val: Value(Scalar(0x0000002a)) }
+          _1 = const 99i32;                // scope 1 at $DIR/mutable_variable.rs:6:5: 6:11
+                                           // ty::Const
+                                           // + ty: i32
+                                           // + val: Value(Scalar(0x00000063))
+                                           // mir::Constant
+-                                          // + span: $DIR/mutable_variable.rs:6:9: 6:11
++                                          // + span: $DIR/mutable_variable.rs:6:5: 6:11
+                                           // + literal: Const { ty: i32, val: Value(Scalar(0x00000063)) }
+          StorageLive(_2);                 // scope 1 at $DIR/mutable_variable.rs:7:9: 7:10
+-         _2 = _1;                         // scope 1 at $DIR/mutable_variable.rs:7:13: 7:14
++         _2 = const 99i32;                // scope 1 at $DIR/mutable_variable.rs:7:13: 7:14
++                                          // ty::Const
++                                          // + ty: i32
++                                          // + val: Value(Scalar(0x00000063))
++                                          // mir::Constant
++                                          // + span: $DIR/mutable_variable.rs:7:13: 7:14
++                                          // + literal: Const { ty: i32, val: Value(Scalar(0x00000063)) }
+          _0 = const ();                   // scope 0 at $DIR/mutable_variable.rs:4:11: 8:2
+                                           // ty::Const
+                                           // + ty: ()
+                                           // + val: Value(Scalar(<ZST>))
+                                           // mir::Constant
+                                           // + span: $DIR/mutable_variable.rs:4:11: 8:2
+                                           // + literal: Const { ty: (), val: Value(Scalar(<ZST>)) }
+          StorageDead(_2);                 // scope 1 at $DIR/mutable_variable.rs:8:1: 8:2
+          StorageDead(_1);                 // scope 0 at $DIR/mutable_variable.rs:8:1: 8:2
+          return;                          // scope 0 at $DIR/mutable_variable.rs:8:2: 8:2
+      }
+  }
+

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

@@ -0,0 +1,8 @@
+// compile-flags: -O
+
+// EMIT_MIR rustc.main.ConstProp.diff
+fn main() {
+    let mut x = (42, 43);
+    x.1 = 99;
+    let y = x;
+}

src/test/mir-opt/const_prop/mutable_variable_aggregate/rustc.main.ConstProp.diff

@@ -0,0 +1,67 @@
+- // MIR for `main` before ConstProp
++ // MIR for `main` after ConstProp
+
+  fn main() -> () {
+      let mut _0: ();                      // return place in scope 0 at $DIR/mutable_variable_aggregate.rs:4:11: 4:11
+      let mut _1: (i32, i32);              // in scope 0 at $DIR/mutable_variable_aggregate.rs:5:9: 5:14
+      scope 1 {
+          debug x => _1;                   // in scope 1 at $DIR/mutable_variable_aggregate.rs:5:9: 5:14
+          let _2: (i32, i32);              // in scope 1 at $DIR/mutable_variable_aggregate.rs:7:9: 7:10
+          scope 2 {
+              debug y => _2;               // in scope 2 at $DIR/mutable_variable_aggregate.rs:7:9: 7:10
+          }
+      }
+
+      bb0: {
+          StorageLive(_1);                 // scope 0 at $DIR/mutable_variable_aggregate.rs:5:9: 5:14
+          _1 = (const 42i32, const 43i32); // scope 0 at $DIR/mutable_variable_aggregate.rs:5:17: 5:25
+                                           // ty::Const
+                                           // + ty: i32
+                                           // + val: Value(Scalar(0x0000002a))
+                                           // mir::Constant
+-                                          // + span: $DIR/mutable_variable_aggregate.rs:5:18: 5:20
++                                          // + span: $DIR/mutable_variable_aggregate.rs:5:17: 5:25
+                                           // + literal: Const { ty: i32, val: Value(Scalar(0x0000002a)) }
+                                           // ty::Const
+                                           // + ty: i32
+                                           // + val: Value(Scalar(0x0000002b))
+                                           // mir::Constant
+-                                          // + span: $DIR/mutable_variable_aggregate.rs:5:22: 5:24
++                                          // + span: $DIR/mutable_variable_aggregate.rs:5:17: 5:25
+                                           // + literal: Const { ty: i32, val: Value(Scalar(0x0000002b)) }
+          (_1.1: i32) = const 99i32;       // scope 1 at $DIR/mutable_variable_aggregate.rs:6:5: 6:13
+                                           // ty::Const
+                                           // + ty: i32
+                                           // + val: Value(Scalar(0x00000063))
+                                           // mir::Constant
+-                                          // + span: $DIR/mutable_variable_aggregate.rs:6:11: 6:13
++                                          // + span: $DIR/mutable_variable_aggregate.rs:6:5: 6:13
+                                           // + literal: Const { ty: i32, val: Value(Scalar(0x00000063)) }
+          StorageLive(_2);                 // scope 1 at $DIR/mutable_variable_aggregate.rs:7:9: 7:10
+-         _2 = _1;                         // scope 1 at $DIR/mutable_variable_aggregate.rs:7:13: 7:14
++         _2 = (const 42i32, const 99i32); // scope 1 at $DIR/mutable_variable_aggregate.rs:7:13: 7:14
++                                          // ty::Const
++                                          // + ty: i32
++                                          // + val: Value(Scalar(0x0000002a))
++                                          // mir::Constant
++                                          // + span: $DIR/mutable_variable_aggregate.rs:7:13: 7:14
++                                          // + literal: Const { ty: i32, val: Value(Scalar(0x0000002a)) }
++                                          // ty::Const
++                                          // + ty: i32
++                                          // + val: Value(Scalar(0x00000063))
++                                          // mir::Constant
++                                          // + span: $DIR/mutable_variable_aggregate.rs:7:13: 7:14
++                                          // + literal: Const { ty: i32, val: Value(Scalar(0x00000063)) }
+          _0 = const ();                   // scope 0 at $DIR/mutable_variable_aggregate.rs:4:11: 8:2
+                                           // ty::Const
+                                           // + ty: ()
+                                           // + val: Value(Scalar(<ZST>))
+                                           // mir::Constant
+                                           // + span: $DIR/mutable_variable_aggregate.rs:4:11: 8:2
+                                           // + literal: Const { ty: (), val: Value(Scalar(<ZST>)) }
+          StorageDead(_2);                 // scope 1 at $DIR/mutable_variable_aggregate.rs:8:1: 8:2
+          StorageDead(_1);                 // scope 0 at $DIR/mutable_variable_aggregate.rs:8:1: 8:2
+          return;                          // scope 0 at $DIR/mutable_variable_aggregate.rs:8:2: 8:2
+      }
+  }
+

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

@@ -0,0 +1,9 @@
+// compile-flags: -O
+
+// EMIT_MIR rustc.main.ConstProp.diff
+fn main() {
+    let mut x = (42, 43);
+    let z = &mut x;
+    z.1 = 99;
+    let y = x;
+}