Auto merge of #75562 - oli-obk:const_prop_no_aggregates, r=wesleywiser

Check that we don't use `Rvalue::Aggregate` after the deaggregator

fixes #75481

r? @wesleywiser

cc @RalfJung (modified the validator)

Author: bors

src/librustc_middle/mir/mod.rs

@@ -73,15 +73,35 @@ impl<'tcx> HasLocalDecls<'tcx> for Body<'tcx> {
 
 /// The various "big phases" that MIR goes through.
 ///
+/// These phases all describe dialects of MIR. Since all MIR uses the same datastructures, the
+/// dialects forbid certain variants or values in certain phases.
+///
+/// Note: Each phase's validation checks all invariants of the *previous* phases' dialects. A phase
+/// that changes the dialect documents what invariants must be upheld *after* that phase finishes.
+///
 /// Warning: ordering of variants is significant.
 #[derive(Copy, Clone, TyEncodable, TyDecodable, Debug, PartialEq, Eq, PartialOrd, Ord)]
 #[derive(HashStable)]
 pub enum MirPhase {
     Build = 0,
+    // FIXME(oli-obk): it's unclear whether we still need this phase (and its corresponding query).
+    // We used to have this for pre-miri MIR based const eval.
     Const = 1,
-    Validated = 2,
-    DropElab = 3,
-    Optimized = 4,
+    /// This phase checks the MIR for promotable elements and takes them out of the main MIR body
+    /// by creating a new MIR body per promoted element. After this phase (and thus the termination
+    /// of the `mir_promoted` query), these promoted elements are available in the `promoted_mir`
+    /// query.
+    ConstPromotion = 2,
+    /// After this phase
+    /// * the only `AggregateKind`s allowed are `Array` and `Generator`,
+    /// * `DropAndReplace` is gone for good
+    /// * `Drop` now uses explicit drop flags visible in the MIR and reaching a `Drop` terminator
+    ///   means that the auto-generated drop glue will be invoked.
+    DropLowering = 3,
+    /// After this phase, generators are explicit state machines (no more `Yield`).
+    /// `AggregateKind::Generator` is gone for good.
+    GeneratorLowering = 4,
+    Optimization = 5,
 }
 
 impl MirPhase {

src/librustc_middle/query/mod.rs

@@ -247,7 +247,7 @@ rustc_queries! {
             desc { |tcx| "elaborating drops for `{}`", tcx.def_path_str(key.did.to_def_id()) }
         }
 
-        query mir_validated(key: ty::WithOptConstParam<LocalDefId>) ->
+        query mir_promoted(key: ty::WithOptConstParam<LocalDefId>) ->
             (
                 &'tcx Steal<mir::Body<'tcx>>,
                 &'tcx Steal<IndexVec<mir::Promoted, mir::Body<'tcx>>>
@@ -281,6 +281,11 @@ rustc_queries! {
             cache_on_disk_if { key.is_local() }
         }
 
+        /// The `DefId` is the `DefId` of the containing MIR body. Promoteds do not have their own
+        /// `DefId`. This function returns all promoteds in the specified body. The body references
+        /// promoteds by the `DefId` and the `mir::Promoted` index. This is necessary, because
+        /// after inlining a body may refer to promoteds from other bodies. In that case you still
+        /// need to use the `DefId` of the original body.
         query promoted_mir(key: DefId) -> &'tcx IndexVec<mir::Promoted, mir::Body<'tcx>> {
             desc { |tcx| "optimizing promoted MIR for `{}`", tcx.def_path_str(key) }
             cache_on_disk_if { key.is_local() }

src/librustc_middle/ty/query/mod.rs

@@ -133,7 +133,7 @@ rustc_query_append! { [define_queries!][<'tcx>] }
 /// `DefPathHash` in the current codebase to the corresponding `DefId`, we have
 /// everything we need to re-run the query.
 ///
-/// Take the `mir_validated` query as an example. Like many other queries, it
+/// Take the `mir_promoted` query as an example. Like many other queries, it
 /// just has a single parameter: the `DefId` of the item it will compute the
 /// validated MIR for. Now, when we call `force_from_dep_node()` on a `DepNode`
 /// with kind `MirValidated`, we know that the GUID/fingerprint of the `DepNode`

src/librustc_mir/borrow_check/mod.rs

@@ -106,7 +106,7 @@ fn mir_borrowck<'tcx>(
     tcx: TyCtxt<'tcx>,
     def: ty::WithOptConstParam<LocalDefId>,
 ) -> &'tcx BorrowCheckResult<'tcx> {
-    let (input_body, promoted) = tcx.mir_validated(def);
+    let (input_body, promoted) = tcx.mir_promoted(def);
     debug!("run query mir_borrowck: {}", tcx.def_path_str(def.did.to_def_id()));
 
     let opt_closure_req = tcx.infer_ctxt().enter(|infcx| {

src/librustc_mir/interpret/intern.rs

@@ -7,7 +7,8 @@ use super::validity::RefTracking;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_hir as hir;
 use rustc_middle::mir::interpret::InterpResult;
-use rustc_middle::ty::{self, query::TyCtxtAt, Ty};
+use rustc_middle::ty::{self, layout::TyAndLayout, query::TyCtxtAt, Ty};
+use rustc_target::abi::Size;
 
 use rustc_ast::Mutability;
 
@@ -430,3 +431,25 @@ pub fn intern_const_alloc_recursive<M: CompileTimeMachine<'mir, 'tcx>>(
         }
     }
 }
+
+impl<'mir, 'tcx: 'mir, M: super::intern::CompileTimeMachine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
+    /// A helper function that allocates memory for the layout given and gives you access to mutate
+    /// it. Once your own mutation code is done, the backing `Allocation` is removed from the
+    /// current `Memory` and returned.
+    pub(crate) fn intern_with_temp_alloc(
+        &mut self,
+        layout: TyAndLayout<'tcx>,
+        f: impl FnOnce(
+            &mut InterpCx<'mir, 'tcx, M>,
+            MPlaceTy<'tcx, M::PointerTag>,
+        ) -> InterpResult<'tcx, ()>,
+    ) -> InterpResult<'tcx, &'tcx Allocation> {
+        let dest = self.allocate(layout, MemoryKind::Stack);
+        f(self, dest)?;
+        let ptr = dest.ptr.assert_ptr();
+        assert_eq!(ptr.offset, Size::ZERO);
+        let mut alloc = self.memory.alloc_map.remove(&ptr.alloc_id).unwrap().1;
+        alloc.mutability = Mutability::Not;
+        Ok(self.tcx.intern_const_alloc(alloc))
+    }
+}

src/librustc_mir/transform/const_prop.rs

@@ -14,9 +14,9 @@ use rustc_middle::mir::visit::{
     MutVisitor, MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor,
 };
 use rustc_middle::mir::{
-    AggregateKind, AssertKind, BasicBlock, BinOp, Body, ClearCrossCrate, Constant, Local,
-    LocalDecl, LocalKind, Location, Operand, Place, Rvalue, SourceInfo, SourceScope,
-    SourceScopeData, Statement, StatementKind, Terminator, TerminatorKind, UnOp, RETURN_PLACE,
+    AssertKind, BasicBlock, BinOp, Body, ClearCrossCrate, Constant, Local, LocalDecl, LocalKind,
+    Location, Operand, Place, Rvalue, SourceInfo, SourceScope, SourceScopeData, Statement,
+    StatementKind, Terminator, TerminatorKind, UnOp, RETURN_PLACE,
 };
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutError, TyAndLayout};
 use rustc_middle::ty::subst::{InternalSubsts, Subst};
@@ -28,9 +28,9 @@ use rustc_trait_selection::traits;
 
 use crate::const_eval::ConstEvalErr;
 use crate::interpret::{
-    self, compile_time_machine, truncate, AllocId, Allocation, Frame, ImmTy, Immediate, InterpCx,
-    LocalState, LocalValue, MemPlace, Memory, MemoryKind, OpTy, Operand as InterpOperand, PlaceTy,
-    Pointer, ScalarMaybeUninit, StackPopCleanup,
+    self, compile_time_machine, truncate, AllocId, Allocation, ConstValue, Frame, ImmTy, Immediate,
+    InterpCx, LocalState, LocalValue, MemPlace, Memory, MemoryKind, OpTy, Operand as InterpOperand,
+    PlaceTy, Pointer, ScalarMaybeUninit, StackPopCleanup,
 };
 use crate::transform::{MirPass, MirSource};
 
@@ -824,44 +824,57 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                     ));
                 }
                 Immediate::ScalarPair(
-                    ScalarMaybeUninit::Scalar(one),
-                    ScalarMaybeUninit::Scalar(two),
+                    ScalarMaybeUninit::Scalar(_),
+                    ScalarMaybeUninit::Scalar(_),
                 ) => {
-                    // 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.
-                    // FIXME: implement the general case stated above ^.
-                    let ty = &value.layout.ty.kind;
+                    // 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(substs) = ty {
+                    if let ty::Tuple(substs) = ty.kind {
                         // Only do it if tuple is also a pair with two scalars
                         if substs.len() == 2 {
-                            let opt_ty1_ty2 = self.use_ecx(|this| {
+                            let alloc = self.use_ecx(|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(|layout| layout.abi.is_scalar())
                                         == Some(true)
                                 };
                                 if ty_is_scalar(ty1) && ty_is_scalar(ty2) {
-                                    Ok(Some((ty1, ty2)))
+                                    let alloc = this
+                                        .ecx
+                                        .intern_with_temp_alloc(value.layout, |ecx, dest| {
+                                            ecx.write_immediate_to_mplace(*imm, dest)
+                                        })
+                                        .unwrap();
+                                    Ok(Some(alloc))
                                 } 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),
-                                    ],
-                                );
+                            if let Some(Some(alloc)) = alloc {
+                                // Assign entire constant in a single statement.
+                                // We can't use aggregates, as we run after the aggregate-lowering `MirPhase`.
+                                *rval = Rvalue::Use(Operand::Constant(Box::new(Constant {
+                                    span: source_info.span,
+                                    user_ty: None,
+                                    literal: self.ecx.tcx.mk_const(ty::Const {
+                                        ty,
+                                        val: ty::ConstKind::Value(ConstValue::ByRef {
+                                            alloc,
+                                            offset: Size::ZERO,
+                                        }),
+                                    }),
+                                })));
                             }
                         }
                     }
                 }
+                // Scalars or scalar pairs that contain undef values are assumed to not have
+                // successfully evaluated and are thus not propagated.
                 _ => {}
             }
         }

src/librustc_mir/transform/generator.rs

@@ -57,6 +57,7 @@ use crate::transform::no_landing_pads::no_landing_pads;
 use crate::transform::simplify;
 use crate::transform::{MirPass, MirSource};
 use crate::util::dump_mir;
+use crate::util::expand_aggregate;
 use crate::util::storage;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
@@ -66,7 +67,7 @@ use rustc_index::bit_set::{BitMatrix, BitSet};
 use rustc_index::vec::{Idx, IndexVec};
 use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
 use rustc_middle::mir::*;
-use rustc_middle::ty::subst::SubstsRef;
+use rustc_middle::ty::subst::{Subst, SubstsRef};
 use rustc_middle::ty::GeneratorSubsts;
 use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt};
 use rustc_target::abi::VariantIdx;
@@ -236,10 +237,28 @@ struct TransformVisitor<'tcx> {
 }
 
 impl TransformVisitor<'tcx> {
-    // Make a GeneratorState rvalue
-    fn make_state(&self, idx: VariantIdx, val: Operand<'tcx>) -> Rvalue<'tcx> {
-        let adt = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
-        Rvalue::Aggregate(box adt, vec![val])
+    // Make a GeneratorState variant assignment. `core::ops::GeneratorState` only has single
+    // element tuple variants, so we can just write to the downcasted first field and then set the
+    // discriminant to the appropriate variant.
+    fn make_state(
+        &self,
+        idx: VariantIdx,
+        val: Operand<'tcx>,
+        source_info: SourceInfo,
+    ) -> impl Iterator<Item = Statement<'tcx>> {
+        let kind = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
+        assert_eq!(self.state_adt_ref.variants[idx].fields.len(), 1);
+        let ty = self
+            .tcx
+            .type_of(self.state_adt_ref.variants[idx].fields[0].did)
+            .subst(self.tcx, self.state_substs);
+        expand_aggregate(
+            Place::return_place(),
+            std::iter::once((val, ty)),
+            kind,
+            source_info,
+            self.tcx,
+        )
     }
 
     // Create a Place referencing a generator struct field
@@ -325,13 +344,7 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
         if let Some((state_idx, resume, v, drop)) = ret_val {
             let source_info = data.terminator().source_info;
             // We must assign the value first in case it gets declared dead below
-            data.statements.push(Statement {
-                source_info,
-                kind: StatementKind::Assign(box (
-                    Place::return_place(),
-                    self.make_state(state_idx, v),
-                )),
-            });
+            data.statements.extend(self.make_state(state_idx, v, source_info));
             let state = if let Some((resume, resume_arg)) = resume {
                 // Yield
                 let state = 3 + self.suspension_points.len();