Auto merge of #107267 - cjgillot:keep-aggregate, r=oli-obk

Do not deaggregate MIR

This turns out to simplify a lot of things.
I haven't checked the consequences for miri yet.

cc `@JakobDegen`
r? `@oli-obk`

compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

@@ -2645,6 +2645,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                             operands,
                         ) = rvalue
                         {
+                            let def_id = def_id.expect_local();
                             for operand in operands {
                                 let (Operand::Copy(assigned_from) | Operand::Move(assigned_from)) = operand else {
                                     continue;
@@ -2667,7 +2668,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                                 // into a place then we should annotate the closure in
                                 // case it ends up being assigned into the return place.
                                 annotated_closure =
-                                    self.annotate_fn_sig(*def_id, substs.as_closure().sig());
+                                    self.annotate_fn_sig(def_id, substs.as_closure().sig());
                                 debug!(
                                     "annotate_argument_and_return_for_borrow: \
                                      annotated_closure={:?} assigned_from_local={:?} \

compiler/rustc_borrowck/src/diagnostics/mod.rs

@@ -817,6 +817,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
             && let AggregateKind::Closure(def_id, _) | AggregateKind::Generator(def_id, _, _) = **kind
         {
             debug!("move_spans: def_id={:?} places={:?}", def_id, places);
+            let def_id = def_id.expect_local();
             if let Some((args_span, generator_kind, capture_kind_span, path_span)) =
                 self.closure_span(def_id, moved_place, places)
             {
@@ -945,6 +946,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     box AggregateKind::Generator(def_id, _, _) => (def_id, true),
                     _ => continue,
                 };
+                let def_id = def_id.expect_local();
 
                 debug!(
                     "borrow_spans: def_id={:?} is_generator={:?} places={:?}",

compiler/rustc_borrowck/src/lib.rs

@@ -1278,6 +1278,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                 // in order to populate our used_mut set.
                 match **aggregate_kind {
                     AggregateKind::Closure(def_id, _) | AggregateKind::Generator(def_id, _, _) => {
+                        let def_id = def_id.expect_local();
                         let BorrowCheckResult { used_mut_upvars, .. } =
                             self.infcx.tcx.mir_borrowck(def_id);
                         debug!("{:?} used_mut_upvars={:?}", def_id, used_mut_upvars);

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -2536,7 +2536,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             // clauses on the struct.
             AggregateKind::Closure(def_id, substs)
             | AggregateKind::Generator(def_id, substs, _) => {
-                (def_id.to_def_id(), self.prove_closure_bounds(tcx, def_id, substs, location))
+                (def_id, self.prove_closure_bounds(tcx, def_id.expect_local(), substs, location))
             }
 
             AggregateKind::Array(_) | AggregateKind::Tuple => {

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -13,6 +13,7 @@ use rustc_middle::ty::cast::{CastTy, IntTy};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
 use rustc_middle::ty::{self, adjustment::PointerCast, Instance, Ty, TyCtxt};
 use rustc_span::source_map::{Span, DUMMY_SP};
+use rustc_target::abi::VariantIdx;
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     #[instrument(level = "trace", skip(self, bx))]
@@ -106,31 +107,31 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             mir::Rvalue::Aggregate(ref kind, ref operands) => {
-                let (dest, active_field_index) = match **kind {
-                    mir::AggregateKind::Adt(adt_did, variant_index, _, _, active_field_index) => {
-                        dest.codegen_set_discr(bx, variant_index);
-                        if bx.tcx().adt_def(adt_did).is_enum() {
-                            (dest.project_downcast(bx, variant_index), active_field_index)
-                        } else {
-                            (dest, active_field_index)
-                        }
+                let (variant_index, variant_dest, active_field_index) = match **kind {
+                    mir::AggregateKind::Adt(_, variant_index, _, _, active_field_index) => {
+                        let variant_dest = dest.project_downcast(bx, variant_index);
+                        (variant_index, variant_dest, active_field_index)
                     }
-                    _ => (dest, None),
+                    _ => (VariantIdx::from_u32(0), dest, None),
                 };
+                if active_field_index.is_some() {
+                    assert_eq!(operands.len(), 1);
+                }
                 for (i, operand) in operands.iter().enumerate() {
                     let op = self.codegen_operand(bx, operand);
                     // Do not generate stores and GEPis for zero-sized fields.
                     if !op.layout.is_zst() {
                         let field_index = active_field_index.unwrap_or(i);
                         let field = if let mir::AggregateKind::Array(_) = **kind {
                             let llindex = bx.cx().const_usize(field_index as u64);
-                            dest.project_index(bx, llindex)
+                            variant_dest.project_index(bx, llindex)
                         } else {
-                            dest.project_field(bx, field_index)
+                            variant_dest.project_field(bx, field_index)
                         };
                         op.val.store(bx, field);
                     }
                 }
+                dest.codegen_set_discr(bx, variant_index);
             }
 
             _ => {

compiler/rustc_const_eval/src/interpret/place.rs

@@ -774,15 +774,6 @@ where
         variant_index: VariantIdx,
         dest: &PlaceTy<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx> {
-        // This must be an enum or generator.
-        match dest.layout.ty.kind() {
-            ty::Adt(adt, _) => assert!(adt.is_enum()),
-            ty::Generator(..) => {}
-            _ => span_bug!(
-                self.cur_span(),
-                "write_discriminant called on non-variant-type (neither enum nor generator)"
-            ),
-        }
         // Layout computation excludes uninhabited variants from consideration
         // therefore there's no way to represent those variants in the given layout.
         // Essentially, uninhabited variants do not have a tag that corresponds to their
@@ -855,6 +846,34 @@ where
         Ok(())
     }
 
+    /// Writes the discriminant of the given variant.
+    #[instrument(skip(self), level = "debug")]
+    pub fn write_aggregate(
+        &mut self,
+        kind: &mir::AggregateKind<'tcx>,
+        operands: &[mir::Operand<'tcx>],
+        dest: &PlaceTy<'tcx, M::Provenance>,
+    ) -> InterpResult<'tcx> {
+        self.write_uninit(&dest)?;
+        let (variant_index, variant_dest, active_field_index) = match *kind {
+            mir::AggregateKind::Adt(_, variant_index, _, _, active_field_index) => {
+                let variant_dest = self.place_downcast(&dest, variant_index)?;
+                (variant_index, variant_dest, active_field_index)
+            }
+            _ => (VariantIdx::from_u32(0), dest.clone(), None),
+        };
+        if active_field_index.is_some() {
+            assert_eq!(operands.len(), 1);
+        }
+        for (field_index, operand) in operands.iter().enumerate() {
+            let field_index = active_field_index.unwrap_or(field_index);
+            let field_dest = self.place_field(&variant_dest, field_index)?;
+            let op = self.eval_operand(operand, Some(field_dest.layout))?;
+            self.copy_op(&op, &field_dest, /*allow_transmute*/ false)?;
+        }
+        self.write_discriminant(variant_index, &dest)
+    }
+
     pub fn raw_const_to_mplace(
         &self,
         raw: ConstAlloc<'tcx>,

compiler/rustc_const_eval/src/interpret/step.rs

@@ -199,13 +199,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
 
             Aggregate(box ref kind, ref operands) => {
-                assert!(matches!(kind, mir::AggregateKind::Array(..)));
-
-                for (field_index, operand) in operands.iter().enumerate() {
-                    let op = self.eval_operand(operand, None)?;
-                    let field_dest = self.place_field(&dest, field_index)?;
-                    self.copy_op(&op, &field_dest, /*allow_transmute*/ false)?;
-                }
+                self.write_aggregate(kind, operands, &dest)?;
             }
 
             Repeat(ref operand, _) => {

compiler/rustc_const_eval/src/transform/check_consts/check.rs

@@ -453,7 +453,7 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
 
             Rvalue::Aggregate(kind, ..) => {
                 if let AggregateKind::Generator(def_id, ..) = kind.as_ref()
-                    && let Some(generator_kind @ hir::GeneratorKind::Async(..)) = self.tcx.generator_kind(def_id.to_def_id())
+                    && let Some(generator_kind @ hir::GeneratorKind::Async(..)) = self.tcx.generator_kind(def_id)
                 {
                     self.check_op(ops::Generator(generator_kind));
                 }

compiler/rustc_const_eval/src/transform/validate.rs

@@ -8,10 +8,10 @@ use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::visit::NonUseContext::VarDebugInfo;
 use rustc_middle::mir::visit::{PlaceContext, Visitor};
 use rustc_middle::mir::{
-    traversal, AggregateKind, BasicBlock, BinOp, Body, BorrowKind, CastKind, CopyNonOverlapping,
-    Local, Location, MirPass, MirPhase, NonDivergingIntrinsic, Operand, Place, PlaceElem, PlaceRef,
-    ProjectionElem, RetagKind, RuntimePhase, Rvalue, SourceScope, Statement, StatementKind,
-    Terminator, TerminatorKind, UnOp, START_BLOCK,
+    traversal, BasicBlock, BinOp, Body, BorrowKind, CastKind, CopyNonOverlapping, Local, Location,
+    MirPass, MirPhase, NonDivergingIntrinsic, Operand, Place, PlaceElem, PlaceRef, ProjectionElem,
+    RetagKind, RuntimePhase, Rvalue, SourceScope, Statement, StatementKind, Terminator,
+    TerminatorKind, UnOp, START_BLOCK,
 };
 use rustc_middle::ty::{self, InstanceDef, ParamEnv, Ty, TyCtxt};
 use rustc_mir_dataflow::impls::MaybeStorageLive;
@@ -423,19 +423,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
             };
         }
         match rvalue {
-            Rvalue::Use(_) | Rvalue::CopyForDeref(_) => {}
-            Rvalue::Aggregate(agg_kind, _) => {
-                let disallowed = match **agg_kind {
-                    AggregateKind::Array(..) => false,
-                    _ => self.mir_phase >= MirPhase::Runtime(RuntimePhase::PostCleanup),
-                };
-                if disallowed {
-                    self.fail(
-                        location,
-                        format!("{:?} have been lowered to field assignments", rvalue),
-                    )
-                }
-            }
+            Rvalue::Use(_) | Rvalue::CopyForDeref(_) | Rvalue::Aggregate(..) => {}
             Rvalue::Ref(_, BorrowKind::Shallow, _) => {
                 if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Initial) {
                     self.fail(

compiler/rustc_const_eval/src/util/mod.rs

@@ -1,4 +1,3 @@
-pub mod aggregate;
 mod alignment;
 mod call_kind;
 pub mod collect_writes;
@@ -7,7 +6,6 @@ mod find_self_call;
 mod might_permit_raw_init;
 mod type_name;
 
-pub use self::aggregate::expand_aggregate;
 pub use self::alignment::is_disaligned;
 pub use self::call_kind::{call_kind, CallDesugaringKind, CallKind};
 pub use self::compare_types::{is_equal_up_to_subtyping, is_subtype};

compiler/rustc_middle/src/mir/mod.rs

@@ -2098,10 +2098,7 @@ impl<'tcx> Debug for Rvalue<'tcx> {
                     AggregateKind::Closure(def_id, substs) => ty::tls::with(|tcx| {
                         let name = if tcx.sess.opts.unstable_opts.span_free_formats {
                             let substs = tcx.lift(substs).unwrap();
-                            format!(
-                                "[closure@{}]",
-                                tcx.def_path_str_with_substs(def_id.to_def_id(), substs),
-                            )
+                            format!("[closure@{}]", tcx.def_path_str_with_substs(def_id, substs),)
                         } else {
                             let span = tcx.def_span(def_id);
                             format!(
@@ -2112,11 +2109,17 @@ impl<'tcx> Debug for Rvalue<'tcx> {
                         let mut struct_fmt = fmt.debug_struct(&name);
 
                         // FIXME(project-rfc-2229#48): This should be a list of capture names/places
-                        if let Some(upvars) = tcx.upvars_mentioned(def_id) {
+                        if let Some(def_id) = def_id.as_local()
+                            && let Some(upvars) = tcx.upvars_mentioned(def_id)
+                        {
                             for (&var_id, place) in iter::zip(upvars.keys(), places) {
                                 let var_name = tcx.hir().name(var_id);
                                 struct_fmt.field(var_name.as_str(), place);
                             }
+                        } else {
+                            for (index, place) in places.iter().enumerate() {
+                                struct_fmt.field(&format!("{index}"), place);
+                            }
                         }
 
                         struct_fmt.finish()
@@ -2127,11 +2130,17 @@ impl<'tcx> Debug for Rvalue<'tcx> {
                         let mut struct_fmt = fmt.debug_struct(&name);
 
                         // FIXME(project-rfc-2229#48): This should be a list of capture names/places
-                        if let Some(upvars) = tcx.upvars_mentioned(def_id) {
+                        if let Some(def_id) = def_id.as_local()
+                            && let Some(upvars) = tcx.upvars_mentioned(def_id)
+                        {
                             for (&var_id, place) in iter::zip(upvars.keys(), places) {
                                 let var_name = tcx.hir().name(var_id);
                                 struct_fmt.field(var_name.as_str(), place);
                             }
+                        } else {
+                            for (index, place) in places.iter().enumerate() {
+                                struct_fmt.field(&format!("{index}"), place);
+                            }
                         }
 
                         struct_fmt.finish()

compiler/rustc_middle/src/mir/syntax.rs

@@ -1203,10 +1203,8 @@ pub enum AggregateKind<'tcx> {
     /// active field index would identity the field `c`
     Adt(DefId, VariantIdx, SubstsRef<'tcx>, Option<UserTypeAnnotationIndex>, Option<usize>),
 
-    // Note: We can use LocalDefId since closures and generators a deaggregated
-    // before codegen.
-    Closure(LocalDefId, SubstsRef<'tcx>),
-    Generator(LocalDefId, SubstsRef<'tcx>, hir::Movability),
+    Closure(DefId, SubstsRef<'tcx>),
+    Generator(DefId, SubstsRef<'tcx>, hir::Movability),
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, TyEncodable, TyDecodable, Hash, HashStable)]

compiler/rustc_middle/src/mir/tcx.rs

@@ -205,9 +205,9 @@ impl<'tcx> Rvalue<'tcx> {
                 AggregateKind::Adt(did, _, substs, _, _) => {
                     tcx.bound_type_of(did).subst(tcx, substs)
                 }
-                AggregateKind::Closure(did, substs) => tcx.mk_closure(did.to_def_id(), substs),
+                AggregateKind::Closure(did, substs) => tcx.mk_closure(did, substs),
                 AggregateKind::Generator(did, substs, movability) => {
-                    tcx.mk_generator(did.to_def_id(), substs, movability)
+                    tcx.mk_generator(did, substs, movability)
                 }
             },
             Rvalue::ShallowInitBox(_, ty) => tcx.mk_box(ty),

compiler/rustc_mir_build/src/build/expr/as_rvalue.rs

@@ -439,10 +439,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         // We implicitly set the discriminant to 0. See
                         // librustc_mir/transform/deaggregator.rs for details.
                         let movability = movability.unwrap();
-                        Box::new(AggregateKind::Generator(closure_id, substs, movability))
+                        Box::new(AggregateKind::Generator(
+                            closure_id.to_def_id(),
+                            substs,
+                            movability,
+                        ))
                     }
                     UpvarSubsts::Closure(substs) => {
-                        Box::new(AggregateKind::Closure(closure_id, substs))
+                        Box::new(AggregateKind::Closure(closure_id.to_def_id(), substs))
                     }
                 };
                 block.and(Rvalue::Aggregate(result, operands))