Auto merge of rust-lang#113858 - cjgillot:const-prop-pairs, r=oli-obk

Always const-prop scalars and scalar pairs

This removes some complexity from the pass.

The limitation to propagate ScalarPairs only for tuple comes from rust-lang#67015, when ScalarPair constant were modeled using `Rvalue::Aggregate`. Nowadays, we use `ConstValue::ByRef`, which does not care about the underlying type.

The justification for not propagating in all cases was perf. This seems not to be a clear cut any more: rust-lang#113858 (comment)

Author: bors

compiler/rustc_mir_transform/src/const_prop.rs

@@ -14,8 +14,7 @@ use rustc_middle::mir::visit::{
 };
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::{LayoutError, LayoutOf, LayoutOfHelpers, TyAndLayout};
-use rustc_middle::ty::GenericArgs;
-use rustc_middle::ty::{self, ConstKind, Instance, ParamEnv, Ty, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, GenericArgs, Instance, ParamEnv, Ty, TyCtxt, TypeVisitableExt};
 use rustc_span::{def_id::DefId, Span, DUMMY_SP};
 use rustc_target::abi::{self, Align, HasDataLayout, Size, TargetDataLayout};
 use rustc_target::spec::abi::Abi as CallAbi;
@@ -407,51 +406,9 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         ecx.machine.written_only_inside_own_block_locals.remove(&local);
     }
 
-    /// Returns the value, if any, of evaluating `c`.
-    fn eval_constant(&mut self, c: &Constant<'tcx>) -> Option<OpTy<'tcx>> {
-        // FIXME we need to revisit this for #67176
-        if c.has_param() {
-            return None;
-        }
-
-        // No span, we don't want errors to be shown.
-        self.ecx.eval_mir_constant(&c.literal, None, None).ok()
-    }
-
-    /// Returns the value, if any, of evaluating `place`.
-    fn eval_place(&mut self, place: Place<'tcx>) -> Option<OpTy<'tcx>> {
-        trace!("eval_place(place={:?})", place);
-        self.ecx.eval_place_to_op(place, None).ok()
-    }
-
-    /// Returns the value, if any, of evaluating `op`. Calls upon `eval_constant`
-    /// or `eval_place`, depending on the variant of `Operand` used.
-    fn eval_operand(&mut self, op: &Operand<'tcx>) -> Option<OpTy<'tcx>> {
-        match *op {
-            Operand::Constant(ref c) => self.eval_constant(c),
-            Operand::Move(place) | Operand::Copy(place) => self.eval_place(place),
-        }
-    }
-
     fn propagate_operand(&mut self, operand: &mut Operand<'tcx>) {
-        match *operand {
-            Operand::Copy(l) | Operand::Move(l) => {
-                if let Some(value) = self.get_const(l) && self.should_const_prop(&value) {
-                    // FIXME(felix91gr): this code only handles `Scalar` cases.
-                    // For now, we're not handling `ScalarPair` cases because
-                    // doing so here would require a lot of code duplication.
-                    // We should hopefully generalize `Operand` handling into a fn,
-                    // and use it to do const-prop here and everywhere else
-                    // where it makes sense.
-                    if let interpret::Operand::Immediate(interpret::Immediate::Scalar(
-                        scalar,
-                    )) = *value
-                    {
-                        *operand = self.operand_from_scalar(scalar, value.layout.ty);
-                    }
-                }
-            }
-            Operand::Constant(_) => (),
+        if let Some(place) = operand.place() && let Some(op) = self.replace_with_const(place) {
+            *operand = op;
         }
     }
 
@@ -579,93 +536,45 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }))
     }
 
-    fn replace_with_const(&mut self, place: Place<'tcx>, rval: &mut Rvalue<'tcx>) {
+    fn replace_with_const(&mut self, place: Place<'tcx>) -> Option<Operand<'tcx>> {
         // This will return None if the above `const_prop` invocation only "wrote" a
         // type whose creation requires no write. E.g. a generator whose initial state
         // consists solely of uninitialized memory (so it doesn't capture any locals).
-        let Some(ref value) = self.get_const(place) else { return };
-        if !self.should_const_prop(value) {
-            return;
-        }
-        trace!("replacing {:?}={:?} with {:?}", place, rval, value);
-
-        if let Rvalue::Use(Operand::Constant(c)) = rval {
-            match c.literal {
-                ConstantKind::Ty(c) if matches!(c.kind(), ConstKind::Unevaluated(..)) => {}
-                _ => {
-                    trace!("skipping replace of Rvalue::Use({:?} because it is already a const", c);
-                    return;
-                }
-            }
+        let value = self.get_const(place)?;
+        if !self.tcx.consider_optimizing(|| format!("ConstantPropagation - {value:?}")) {
+            return None;
         }
+        trace!("replacing {:?} with {:?}", place, value);
 
-        trace!("attempting to replace {:?} with {:?}", rval, value);
         // FIXME> figure out what to do when read_immediate_raw fails
-        let imm = self.ecx.read_immediate_raw(value).ok();
+        let imm = self.ecx.read_immediate_raw(&value).ok()?;
 
-        if let Some(Right(imm)) = imm {
-            match *imm {
-                interpret::Immediate::Scalar(scalar) => {
-                    *rval = Rvalue::Use(self.operand_from_scalar(scalar, value.layout.ty));
-                }
-                Immediate::ScalarPair(..) => {
-                    // Found a value represented as a pair. For now only do const-prop if the type
-                    // of `rvalue` is also a tuple with two scalars.
-                    // FIXME: enable the general case stated above ^.
-                    let ty = value.layout.ty;
-                    // Only do it for tuples
-                    if let ty::Tuple(types) = ty.kind() {
-                        // Only do it if tuple is also a pair with two scalars
-                        if let [ty1, ty2] = types[..] {
-                            let ty_is_scalar = |ty| {
-                                self.ecx.layout_of(ty).ok().map(|layout| layout.abi.is_scalar())
-                                    == Some(true)
-                            };
-                            let alloc = if ty_is_scalar(ty1) && ty_is_scalar(ty2) {
-                                let alloc = self
-                                    .ecx
-                                    .intern_with_temp_alloc(value.layout, |ecx, dest| {
-                                        ecx.write_immediate(*imm, dest)
-                                    })
-                                    .unwrap();
-                                Some(alloc)
-                            } else {
-                                None
-                            };
-
-                            if let Some(alloc) = alloc {
-                                // Assign entire constant in a single statement.
-                                // We can't use aggregates, as we run after the aggregate-lowering `MirPhase`.
-                                let const_val = ConstValue::ByRef { alloc, offset: Size::ZERO };
-                                let literal = ConstantKind::Val(const_val, ty);
-                                *rval = Rvalue::Use(Operand::Constant(Box::new(Constant {
-                                    span: DUMMY_SP,
-                                    user_ty: None,
-                                    literal,
-                                })));
-                            }
-                        }
-                    }
-                }
-                // Scalars or scalar pairs that contain undef values are assumed to not have
-                // successfully evaluated and are thus not propagated.
-                _ => {}
+        let Right(imm) = imm else { return None };
+        match *imm {
+            Immediate::Scalar(scalar) if scalar.try_to_int().is_ok() => {
+                Some(self.operand_from_scalar(scalar, value.layout.ty))
             }
-        }
-    }
-
-    /// Returns `true` if and only if this `op` should be const-propagated into.
-    fn should_const_prop(&mut self, op: &OpTy<'tcx>) -> bool {
-        if !self.tcx.consider_optimizing(|| format!("ConstantPropagation - OpTy: {:?}", op)) {
-            return false;
-        }
-
-        match **op {
-            interpret::Operand::Immediate(Immediate::Scalar(s)) => s.try_to_int().is_ok(),
-            interpret::Operand::Immediate(Immediate::ScalarPair(l, r)) => {
-                l.try_to_int().is_ok() && r.try_to_int().is_ok()
+            Immediate::ScalarPair(l, r) if l.try_to_int().is_ok() && r.try_to_int().is_ok() => {
+                let alloc = self
+                    .ecx
+                    .intern_with_temp_alloc(value.layout, |ecx, dest| {
+                        ecx.write_immediate(*imm, dest)
+                    })
+                    .ok()?;
+
+                let literal = ConstantKind::Val(
+                    ConstValue::ByRef { alloc, offset: Size::ZERO },
+                    value.layout.ty,
+                );
+                Some(Operand::Constant(Box::new(Constant {
+                    span: DUMMY_SP,
+                    user_ty: None,
+                    literal,
+                })))
             }
-            _ => false,
+            // Scalars or scalar pairs that contain undef values are assumed to not have
+            // successfully evaluated and are thus not propagated.
+            _ => None,
         }
     }
 
@@ -810,12 +719,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
 
     fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
         self.super_operand(operand, location);
-
-        // Only const prop copies and moves on `mir_opt_level=3` as doing so
-        // currently slightly increases compile time in some cases.
-        if self.tcx.sess.mir_opt_level() >= 3 {
-            self.propagate_operand(operand)
-        }
+        self.propagate_operand(operand)
     }
 
     fn process_projection_elem(
@@ -825,8 +729,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
     ) -> Option<PlaceElem<'tcx>> {
         if let PlaceElem::Index(local) = elem
             && let Some(value) = self.get_const(local.into())
-            && self.should_const_prop(&value)
-            && let interpret::Operand::Immediate(interpret::Immediate::Scalar(scalar)) = *value
+            && let interpret::Operand::Immediate(Immediate::Scalar(scalar)) = *value
             && let Ok(offset) = scalar.to_target_usize(&self.tcx)
             && let Some(min_length) = offset.checked_add(1)
         {
@@ -852,7 +755,14 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
             ConstPropMode::NoPropagation => self.ensure_not_propagated(place.local),
             ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => {
                 if let Some(()) = self.eval_rvalue_with_identities(rvalue, *place) {
-                    self.replace_with_const(*place, rvalue);
+                    // If this was already an evaluated constant, keep it.
+                    if let Rvalue::Use(Operand::Constant(c)) = rvalue
+                        && let ConstantKind::Val(..) = c.literal
+                    {
+                        trace!("skipping replace of Rvalue::Use({:?} because it is already a const", c);
+                    } else if let Some(operand) = self.replace_with_const(*place) {
+                        *rvalue = Rvalue::Use(operand);
+                    }
                 } else {
                     // Const prop failed, so erase the destination, ensuring that whatever happens
                     // from here on, does not know about the previous value.
@@ -919,45 +829,6 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         }
     }
 
-    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
-        self.super_terminator(terminator, location);
-
-        match &mut terminator.kind {
-            TerminatorKind::Assert { expected, ref mut cond, .. } => {
-                if let Some(ref value) = self.eval_operand(&cond)
-                    && let Ok(value_const) = self.ecx.read_scalar(&value)
-                    && self.should_const_prop(value)
-                {
-                    trace!("assertion on {:?} should be {:?}", value, expected);
-                    *cond = self.operand_from_scalar(value_const, self.tcx.types.bool);
-                }
-            }
-            TerminatorKind::SwitchInt { ref mut discr, .. } => {
-                // FIXME: This is currently redundant with `visit_operand`, but sadly
-                // always visiting operands currently causes a perf regression in LLVM codegen, so
-                // `visit_operand` currently only runs for propagates places for `mir_opt_level=4`.
-                self.propagate_operand(discr)
-            }
-            // None of these have Operands to const-propagate.
-            TerminatorKind::Goto { .. }
-            | TerminatorKind::Resume
-            | TerminatorKind::Terminate
-            | TerminatorKind::Return
-            | TerminatorKind::Unreachable
-            | TerminatorKind::Drop { .. }
-            | TerminatorKind::Yield { .. }
-            | TerminatorKind::GeneratorDrop
-            | TerminatorKind::FalseEdge { .. }
-            | TerminatorKind::FalseUnwind { .. }
-            | TerminatorKind::InlineAsm { .. } => {}
-            // Every argument in our function calls have already been propagated in `visit_operand`.
-            //
-            // NOTE: because LLVM codegen gives slight performance regressions with it, so this is
-            // gated on `mir_opt_level=3`.
-            TerminatorKind::Call { .. } => {}
-        }
-    }
-
     fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
         self.super_basic_block_data(block, data);
 

tests/mir-opt/const_prop/aggregate.main.ConstProp.panic-abort.diff

@@ -26,8 +26,9 @@
           StorageLive(_4);
           StorageLive(_5);
 -         _5 = _1;
+-         _4 = foo(move _5) -> [return: bb1, unwind unreachable];
 +         _5 = const 1_u8;
-          _4 = foo(move _5) -> [return: bb1, unwind unreachable];
++         _4 = foo(const 1_u8) -> [return: bb1, unwind unreachable];
       }
 
       bb1: {

tests/mir-opt/const_prop/aggregate.main.ConstProp.panic-unwind.diff

@@ -26,8 +26,9 @@
           StorageLive(_4);
           StorageLive(_5);
 -         _5 = _1;
+-         _4 = foo(move _5) -> [return: bb1, unwind continue];
 +         _5 = const 1_u8;
-          _4 = foo(move _5) -> [return: bb1, unwind continue];
++         _4 = foo(const 1_u8) -> [return: bb1, unwind continue];
       }
 
       bb1: {

tests/mir-opt/const_prop/aggregate.main.PreCodegen.after.panic-abort.mir

@@ -23,7 +23,7 @@ fn main() -> () {
         StorageLive(_4);
         StorageLive(_5);
         _5 = const 1_u8;
-        _4 = foo(move _5) -> [return: bb1, unwind unreachable];
+        _4 = foo(const 1_u8) -> [return: bb1, unwind unreachable];
     }
 
     bb1: {

tests/mir-opt/const_prop/aggregate.main.PreCodegen.after.panic-unwind.mir

@@ -23,7 +23,7 @@ fn main() -> () {
         StorageLive(_4);
         StorageLive(_5);
         _5 = const 1_u8;
-        _4 = foo(move _5) -> [return: bb1, unwind continue];
+        _4 = foo(const 1_u8) -> [return: bb1, unwind continue];
     }
 
     bb1: {

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

@@ -23,7 +23,7 @@
 -         assert(move _5, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind unreachable];
 +         _4 = const 4_usize;
 +         _5 = const true;
-+         assert(const true, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind unreachable];
++         assert(const true, "index out of bounds: the length is {} but the index is {}", const 4_usize, const 2_usize) -> [success: bb1, unwind unreachable];
       }
 
       bb1: {

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

@@ -23,7 +23,7 @@
 -         assert(move _5, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind continue];
 +         _4 = const 4_usize;
 +         _5 = const true;
-+         assert(const true, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind continue];
++         assert(const true, "index out of bounds: the length is {} but the index is {}", const 4_usize, const 2_usize) -> [success: bb1, unwind continue];
       }
 
       bb1: {

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

@@ -23,7 +23,7 @@
 -         assert(move _5, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind unreachable];
 +         _4 = const 4_usize;
 +         _5 = const true;
-+         assert(const true, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind unreachable];
++         assert(const true, "index out of bounds: the length is {} but the index is {}", const 4_usize, const 2_usize) -> [success: bb1, unwind unreachable];
       }
 
       bb1: {

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

@@ -23,7 +23,7 @@
 -         assert(move _5, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind continue];
 +         _4 = const 4_usize;
 +         _5 = const true;
-+         assert(const true, "index out of bounds: the length is {} but the index is {}", move _4, _3) -> [success: bb1, unwind continue];
++         assert(const true, "index out of bounds: the length is {} but the index is {}", const 4_usize, const 2_usize) -> [success: bb1, unwind continue];
       }
 
       bb1: {

tests/mir-opt/const_prop/bad_op_div_by_zero.main.ConstProp.panic-abort.diff

@@ -38,11 +38,12 @@
 +         _5 = const false;
 +         _6 = const false;
 +         _7 = const false;
-+         assert(!const false, "attempt to compute `{} / {}`, which would overflow", const 1_i32, _3) -> [success: bb2, unwind unreachable];
++         assert(!const false, "attempt to compute `{} / {}`, which would overflow", const 1_i32, const 0_i32) -> [success: bb2, unwind unreachable];
       }
 
       bb2: {
-          _2 = Div(const 1_i32, move _3);
+-         _2 = Div(const 1_i32, move _3);
++         _2 = Div(const 1_i32, const 0_i32);
           StorageDead(_3);
           _0 = const ();
           StorageDead(_2);