Auto merge of #67130 - wesleywiser:const_prop_into_locals, r=<try>

Const prop should finish propagation into user defined variables

Fixes #66638

Temporarily rebased on top of #67015 to get those fixes.

r? @oli-obk

Author: bors

Diff for: 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.
@@ -636,19 +636,45 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                     ScalarMaybeUndef::Scalar(one),
                     ScalarMaybeUndef::Scalar(two)
                 ) => {
+                    // Found a value represented as a pair. For now only do cont-prop if type of
+                    // Rvalue is also a pair with two scalars. The more general case is more
+                    // complicated to implement so we'll do it later.
                     let ty = &value.layout.ty.kind;
+                    // Only do it for tuples
                     if let ty::Tuple(substs) = ty {
-                        *rval = Rvalue::Aggregate(
-                            Box::new(AggregateKind::Tuple),
-                            vec![
-                                self.operand_from_scalar(
-                                    one, substs[0].expect_ty(), source_info.span
-                                ),
-                                self.operand_from_scalar(
-                                    two, substs[1].expect_ty(), source_info.span
-                                ),
-                            ],
-                        );
+                        // Only do it if tuple is also a pair with two scalars
+                        if substs.len() == 2 {
+                            let opt_ty1_ty2 = self.use_ecx(source_info, |this| {
+                                let ty1 = substs[0].expect_ty();
+                                let ty2 = substs[1].expect_ty();
+                                let ty_is_scalar = |ty| {
+                                    this.ecx
+                                        .layout_of(ty)
+                                        .ok()
+                                        .map(|ty| ty.details.abi.is_scalar())
+                                        == Some(true)
+                                };
+                                if ty_is_scalar(ty1) && ty_is_scalar(ty2) {
+                                    Ok(Some((ty1, ty2)))
+                                } else {
+                                    Ok(None)
+                                }
+                            });
+
+                            if let Some(Some((ty1, ty2))) = opt_ty1_ty2 {
+                                *rval = Rvalue::Aggregate(
+                                    Box::new(AggregateKind::Tuple),
+                                    vec![
+                                        self.operand_from_scalar(
+                                            one, ty1, source_info.span
+                                        ),
+                                        self.operand_from_scalar(
+                                            two, ty2, source_info.span
+                                        ),
+                                    ],
+                                );
+                            }
+                        }
                     }
                 },
                 _ => { }
@@ -682,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: ReadOnlyBodyCache<'_, '_>) -> IndexVec<Local, bool> {
+    fn check(body: ReadOnlyBodyCache<'_, '_>) -> 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() {
@@ -702,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);
@@ -727,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
             },
@@ -740,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;
             },
         }
     }
@@ -774,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);
@@ -786,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 {

Diff for: src/librustc_target/abi/mod.rs

@@ -802,6 +802,14 @@ impl Abi {
             _ => false,
         }
     }
+
+    /// Returns `true` is this is a scalar type
+    pub fn is_scalar(&self) -> bool {
+        match *self {
+            Abi::Scalar(_) => true,
+            _ => false,
+        }
+    }
 }
 
 rustc_index::newtype_index! {

Diff for: 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

Diff for: 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;
 //  }

Diff for: src/test/mir-opt/const_prop/issue-66971.rs

@@ -0,0 +1,38 @@
+// compile-flags: -Z mir-opt-level=2
+
+// Due to a bug in propagating scalar pairs the assertion below used to fail. In the expected
+// outputs below, after ConstProp this is how _2 would look like with the bug:
+//
+//     _2 = (const Scalar(0x00) : (), const 0u8);
+//
+// Which has the wrong type.
+
+fn encode(this: ((), u8, u8)) {
+    assert!(this.2 == 0);
+}
+
+fn main() {
+    encode(((), 0, 0));
+}
+
+// END RUST SOURCE
+// START rustc.main.ConstProp.before.mir
+//  bb0: {
+//      ...
+//      _3 = ();
+//      _2 = (move _3, const 0u8, const 0u8);
+//      ...
+//      _1 = const encode(move _2) -> bb1;
+//      ...
+//  }
+// END rustc.main.ConstProp.before.mir
+// START rustc.main.ConstProp.after.mir
+//  bb0: {
+//      ...
+//      _3 = const Scalar(<ZST>) : ();
+//      _2 = (move _3, const 0u8, const 0u8);
+//      ...
+//      _1 = const encode(move _2) -> bb1;
+//      ...
+//  }
+// END rustc.main.ConstProp.after.mir

Diff for: src/test/mir-opt/const_prop/issue-67019.rs

@@ -0,0 +1,34 @@
+// compile-flags: -Z mir-opt-level=2
+
+// This used to ICE in const-prop
+
+fn test(this: ((u8, u8),)) {
+    assert!((this.0).0 == 1);
+}
+
+fn main() {
+    test(((1, 2),));
+}
+
+// Important bit is parameter passing so we only check that below
+// END RUST SOURCE
+// START rustc.main.ConstProp.before.mir
+//  bb0: {
+//      ...
+//      _3 = (const 1u8, const 2u8);
+//      _2 = (move _3,);
+//      ...
+//      _1 = const test(move _2) -> bb1;
+//      ...
+//  }
+// END rustc.main.ConstProp.before.mir
+// START rustc.main.ConstProp.after.mir
+//  bb0: {
+//      ...
+//      _3 = (const 1u8, const 2u8);
+//      _2 = (move _3,);
+//      ...
+//      _1 = const test(move _2) -> bb1;
+//      ...
+//  }
+// END rustc.main.ConstProp.after.mir