Auto merge of #116270 - cjgillot:gvn-aggregate, r=oli-obk,RalfJung

See through aggregates in GVN

This PR is extracted from #111344

The first 2 commit are cleanups to avoid repeated work. I propose to stop removing useless assignments as part of this pass, and let a later `SimplifyLocals` do it. This makes tests easier to read (among others).

The next 3 commits add a constant folding mechanism to the GVN pass, presented in #116012. ~This pass is designed to only use global allocations, to avoid any risk of accidental modification of the stored state.~

The following commits implement opportunistic simplifications, in particular:
- projections of aggregates: `MyStruct { x: a }.x` gets replaced by `a`, works with enums too;
- projections of arrays: `[a, b][0]` becomes `a`;
- projections of repeat expressions: `[a; N][x]` becomes `a`;
- transform arrays of equal operands into a repeat rvalue.

Fixes rust-lang/miri#3090

r? `@oli-obk`

compiler/rustc_const_eval/src/interpret/discriminant.rs

@@ -1,7 +1,8 @@
 //! Functions for reading and writing discriminants of multi-variant layouts (enums and coroutines).
 
-use rustc_middle::ty::layout::{LayoutOf, PrimitiveExt, TyAndLayout};
-use rustc_middle::{mir, ty};
+use rustc_middle::mir;
+use rustc_middle::ty::layout::{LayoutOf, PrimitiveExt};
+use rustc_middle::ty::{self, Ty};
 use rustc_target::abi::{self, TagEncoding};
 use rustc_target::abi::{VariantIdx, Variants};
 
@@ -244,11 +245,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 
     pub fn discriminant_for_variant(
         &self,
-        layout: TyAndLayout<'tcx>,
+        ty: Ty<'tcx>,
         variant: VariantIdx,
     ) -> InterpResult<'tcx, ImmTy<'tcx, M::Provenance>> {
-        let discr_layout = self.layout_of(layout.ty.discriminant_ty(*self.tcx))?;
-        let discr_value = match layout.ty.discriminant_for_variant(*self.tcx, variant) {
+        let discr_layout = self.layout_of(ty.discriminant_ty(*self.tcx))?;
+        let discr_value = match ty.discriminant_for_variant(*self.tcx, variant) {
             Some(discr) => {
                 // This type actually has discriminants.
                 assert_eq!(discr.ty, discr_layout.ty);

compiler/rustc_const_eval/src/interpret/intern.rs

@@ -450,6 +450,42 @@ pub fn intern_const_alloc_recursive<
     Ok(())
 }
 
+/// Intern `ret`. This function assumes that `ret` references no other allocation.
+#[instrument(level = "debug", skip(ecx))]
+pub fn intern_const_alloc_for_constprop<
+    'mir,
+    'tcx: 'mir,
+    T,
+    M: CompileTimeMachine<'mir, 'tcx, T>,
+>(
+    ecx: &mut InterpCx<'mir, 'tcx, M>,
+    alloc_id: AllocId,
+) -> InterpResult<'tcx, ()> {
+    // Move allocation to `tcx`.
+    let Some((_, mut alloc)) = ecx.memory.alloc_map.remove(&alloc_id) else {
+        // Pointer not found in local memory map. It is either a pointer to the global
+        // map, or dangling.
+        if ecx.tcx.try_get_global_alloc(alloc_id).is_none() {
+            throw_ub!(DeadLocal)
+        }
+        // The constant is already in global memory. Do nothing.
+        return Ok(());
+    };
+
+    alloc.mutability = Mutability::Not;
+
+    // We are not doing recursive interning, so we don't currently support provenance.
+    // (If this assertion ever triggers, we should just implement a
+    // proper recursive interning loop.)
+    assert!(alloc.provenance().ptrs().is_empty());
+
+    // Link the alloc id to the actual allocation
+    let alloc = ecx.tcx.mk_const_alloc(alloc);
+    ecx.tcx.set_alloc_id_memory(alloc_id, alloc);
+
+    Ok(())
+}
+
 impl<'mir, 'tcx: 'mir, M: super::intern::CompileTimeMachine<'mir, 'tcx, !>>
     InterpCx<'mir, 'tcx, M>
 {

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -218,7 +218,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             sym::discriminant_value => {
                 let place = self.deref_pointer(&args[0])?;
                 let variant = self.read_discriminant(&place)?;
-                let discr = self.discriminant_for_variant(place.layout, variant)?;
+                let discr = self.discriminant_for_variant(place.layout.ty, variant)?;
                 self.write_immediate(*discr, dest)?;
             }
             sym::exact_div => {

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -1011,7 +1011,7 @@ impl<'tcx, 'a, Prov: Provenance, Extra, Bytes: AllocBytes> AllocRef<'a, 'tcx, Pr
     }
 
     /// Returns whether the allocation has provenance anywhere in the range of the `AllocRef`.
-    pub(crate) fn has_provenance(&self) -> bool {
+    pub fn has_provenance(&self) -> bool {
         !self.alloc.provenance().range_empty(self.range, &self.tcx)
     }
 }

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -21,7 +21,9 @@ mod visitor;
 pub use rustc_middle::mir::interpret::*; // have all the `interpret` symbols in one place: here
 
 pub use self::eval_context::{Frame, FrameInfo, InterpCx, StackPopCleanup};
-pub use self::intern::{intern_const_alloc_recursive, InternKind};
+pub use self::intern::{
+    intern_const_alloc_for_constprop, intern_const_alloc_recursive, InternKind,
+};
 pub use self::machine::{compile_time_machine, AllocMap, Machine, MayLeak, StackPopJump};
 pub use self::memory::{AllocKind, AllocRef, AllocRefMut, FnVal, Memory, MemoryKind};
 pub use self::operand::{ImmTy, Immediate, OpTy, Readable};

compiler/rustc_const_eval/src/interpret/operand.rs

@@ -169,6 +169,16 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
         ImmTy { imm: val.into(), layout }
     }
 
+    #[inline]
+    pub fn from_scalar_pair(a: Scalar<Prov>, b: Scalar<Prov>, layout: TyAndLayout<'tcx>) -> Self {
+        debug_assert!(
+            matches!(layout.abi, Abi::ScalarPair(..)),
+            "`ImmTy::from_scalar_pair` on non-scalar-pair layout"
+        );
+        let imm = Immediate::ScalarPair(a, b);
+        ImmTy { imm, layout }
+    }
+
     #[inline(always)]
     pub fn from_immediate(imm: Immediate<Prov>, layout: TyAndLayout<'tcx>) -> Self {
         debug_assert!(

compiler/rustc_const_eval/src/interpret/step.rs

@@ -297,7 +297,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             Discriminant(place) => {
                 let op = self.eval_place_to_op(place, None)?;
                 let variant = self.read_discriminant(&op)?;
-                let discr = self.discriminant_for_variant(op.layout, variant)?;
+                let discr = self.discriminant_for_variant(op.layout.ty, variant)?;
                 self.write_immediate(*discr, &dest)?;
             }
         }

compiler/rustc_middle/src/mir/consts.rs

@@ -172,6 +172,24 @@ impl<'tcx> ConstValue<'tcx> {
         let end = end.try_into().unwrap();
         Some(data.inner().inspect_with_uninit_and_ptr_outside_interpreter(start..end))
     }
+
+    /// Check if a constant may contain provenance information. This is used by MIR opts.
+    /// Can return `true` even if there is no provenance.
+    pub fn may_have_provenance(&self, tcx: TyCtxt<'tcx>, size: Size) -> bool {
+        match *self {
+            ConstValue::ZeroSized | ConstValue::Scalar(Scalar::Int(_)) => return false,
+            ConstValue::Scalar(Scalar::Ptr(..)) => return true,
+            // It's hard to find out the part of the allocation we point to;
+            // just conservatively check everything.
+            ConstValue::Slice { data, meta: _ } => !data.inner().provenance().ptrs().is_empty(),
+            ConstValue::Indirect { alloc_id, offset } => !tcx
+                .global_alloc(alloc_id)
+                .unwrap_memory()
+                .inner()
+                .provenance()
+                .range_empty(super::AllocRange::from(offset..offset + size), &tcx),
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
@@ -485,6 +503,38 @@ impl<'tcx> Const<'tcx> {
             _ => Self::Ty(c),
         }
     }
+
+    /// Return true if any evaluation of this constant always returns the same value,
+    /// taking into account even pointer identity tests.
+    pub fn is_deterministic(&self) -> bool {
+        // Some constants may generate fresh allocations for pointers they contain,
+        // so using the same constant twice can yield two different results:
+        // - valtrees purposefully generate new allocations
+        // - ConstValue::Slice also generate new allocations
+        match self {
+            Const::Ty(c) => match c.kind() {
+                ty::ConstKind::Param(..) => true,
+                // A valtree may be a reference. Valtree references correspond to a
+                // different allocation each time they are evaluated. Valtrees for primitive
+                // types are fine though.
+                ty::ConstKind::Value(_) => c.ty().is_primitive(),
+                ty::ConstKind::Unevaluated(..) | ty::ConstKind::Expr(..) => false,
+                // Should not appear in runtime MIR.
+                ty::ConstKind::Infer(..)
+                | ty::ConstKind::Bound(..)
+                | ty::ConstKind::Placeholder(..)
+                | ty::ConstKind::Error(..) => bug!(),
+            },
+            Const::Unevaluated(..) => false,
+            // If the same slice appears twice in the MIR, we cannot guarantee that we will
+            // give the same `AllocId` to the data.
+            Const::Val(ConstValue::Slice { .. }, _) => false,
+            Const::Val(
+                ConstValue::ZeroSized | ConstValue::Scalar(_) | ConstValue::Indirect { .. },
+                _,
+            ) => true,
+        }
+    }
 }
 
 /// An unevaluated (potentially generic) constant used in MIR.

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -406,7 +406,8 @@ impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
                 TrackElem::Variant(idx) => self.ecx.project_downcast(op, idx).ok(),
                 TrackElem::Discriminant => {
                     let variant = self.ecx.read_discriminant(op).ok()?;
-                    let discr_value = self.ecx.discriminant_for_variant(op.layout, variant).ok()?;
+                    let discr_value =
+                        self.ecx.discriminant_for_variant(op.layout.ty, variant).ok()?;
                     Some(discr_value.into())
                 }
                 TrackElem::DerefLen => {
@@ -507,7 +508,8 @@ impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
             return None;
         }
         let enum_ty_layout = self.tcx.layout_of(self.param_env.and(enum_ty)).ok()?;
-        let discr_value = self.ecx.discriminant_for_variant(enum_ty_layout, variant_index).ok()?;
+        let discr_value =
+            self.ecx.discriminant_for_variant(enum_ty_layout.ty, variant_index).ok()?;
         Some(discr_value.to_scalar())
     }
 
@@ -854,7 +856,7 @@ impl<'tcx> Visitor<'tcx> for OperandCollector<'tcx, '_, '_, '_> {
     }
 }
 
-struct DummyMachine;
+pub(crate) struct DummyMachine;
 
 impl<'mir, 'tcx: 'mir> rustc_const_eval::interpret::Machine<'mir, 'tcx> for DummyMachine {
     rustc_const_eval::interpret::compile_time_machine!(<'mir, 'tcx>);