Unify `Rvalue::Aggregate` paths in cg_ssa

4 files changed, 31 insertions, 23 deletions

diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index a8d019c9804..ea805c5ad5d 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -1255,7 +1255,7 @@ impl<FieldIdx: Idx> FieldsShape<FieldIdx> {
 
     /// Gets source indices of the fields by increasing offsets.
     #[inline]
-    pub fn index_by_increasing_offset(&self) -> impl Iterator<Item = usize> + '_ {
+    pub fn index_by_increasing_offset(&self) -> impl ExactSizeIterator<Item = usize> + '_ {
         let mut inverse_small = [0u8; 64];
         let mut inverse_big = IndexVec::new();
         let use_small = self.count() <= inverse_small.len();
diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs
index e1c584e8ed5..32fd9b657f9 100644
--- a/compiler/rustc_codegen_ssa/src/mir/operand.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/operand.rs
@@ -110,6 +110,19 @@ impl<V: CodegenObject> OperandValue<V> {
         let (llval, llextra) = self.pointer_parts();
         PlaceValue { llval, llextra, align }
     }
+
+    pub(crate) fn is_expected_variant_for_type<'tcx, Cx: LayoutTypeMethods<'tcx>>(
+        &self,
+        cx: &Cx,
+        ty: TyAndLayout<'tcx>,
+    ) -> bool {
+        match self {
+            OperandValue::ZeroSized => ty.is_zst(),
+            OperandValue::Immediate(_) => cx.is_backend_immediate(ty),
+            OperandValue::Pair(_, _) => cx.is_backend_scalar_pair(ty),
+            OperandValue::Ref(_) => cx.is_backend_ref(ty),
+        }
+    }
 }
 
 /// An `OperandRef` is an "SSA" reference to a Rust value, along with
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index e47a8198736..b764ea6a4c6 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -18,6 +18,7 @@ use rustc_span::{Span, DUMMY_SP};
 use rustc_target::abi::{self, FieldIdx, FIRST_VARIANT};
 
 use arrayvec::ArrayVec;
+use either::Either;
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     #[instrument(level = "trace", skip(self, bx))]
@@ -696,35 +697,25 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 OperandRef { val: OperandValue::Immediate(static_), layout }
             }
             mir::Rvalue::Use(ref operand) => self.codegen_operand(bx, operand),
-            mir::Rvalue::Aggregate(box mir::AggregateKind::RawPtr(..), ref fields) => {
-                let ty = rvalue.ty(self.mir, self.cx.tcx());
-                let layout = self.cx.layout_of(self.monomorphize(ty));
-                let [data, meta] = &*fields.raw else {
-                    bug!("RawPtr fields: {fields:?}");
-                };
-                let data = self.codegen_operand(bx, data);
-                let meta = self.codegen_operand(bx, meta);
-                match (data.val, meta.val) {
-                    (p @ OperandValue::Immediate(_), OperandValue::ZeroSized) => {
-                        OperandRef { val: p, layout }
-                    }
-                    (OperandValue::Immediate(p), OperandValue::Immediate(m)) => {
-                        OperandRef { val: OperandValue::Pair(p, m), layout }
-                    }
-                    _ => bug!("RawPtr operands {data:?} {meta:?}"),
-                }
-            }
             mir::Rvalue::Repeat(..) => bug!("{rvalue:?} in codegen_rvalue_operand"),
-            mir::Rvalue::Aggregate(_, ref fields) => {
+            mir::Rvalue::Aggregate(ref kind, ref fields) => {
                 let ty = rvalue.ty(self.mir, self.cx.tcx());
                 let ty = self.monomorphize(ty);
                 let layout = self.cx.layout_of(ty);
 
+                let field_indices = if let mir::AggregateKind::RawPtr(..) = **kind {
+                    // `index_by_increasing_offset` gives an empty iterator for primitives
+                    Either::Left([0_usize, 1_usize].iter().copied())
+                } else {
+                    Either::Right(layout.fields.index_by_increasing_offset())
+                };
+                debug_assert_eq!(field_indices.len(), fields.len());
+
                 // `rvalue_creates_operand` has arranged that we only get here if
                 // we can build the aggregate immediate from the field immediates.
                 let mut inputs = ArrayVec::<Bx::Value, 2>::new();
                 let mut input_scalars = ArrayVec::<abi::Scalar, 2>::new();
-                for field_idx in layout.fields.index_by_increasing_offset() {
+                for field_idx in field_indices {
                     let field_idx = FieldIdx::from_usize(field_idx);
                     let op = self.codegen_operand(bx, &fields[field_idx]);
                     let values = op.val.immediates_or_place().left_or_else(|p| {
@@ -748,6 +739,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 );
 
                 let val = OperandValue::from_immediates(inputs);
+                debug_assert!(
+                    val.is_expected_variant_for_type(self.cx, layout),
+                    "Made wrong variant {val:?} for type {layout:?}",
+                );
                 OperandRef { val, layout }
             }
             mir::Rvalue::ShallowInitBox(ref operand, content_ty) => {
@@ -792,7 +787,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         debug_assert!(
             if bx.cx().type_has_metadata(ty) {
                 matches!(val, OperandValue::Pair(..))
-        } else {
+            } else {
                 matches!(val, OperandValue::Immediate(..))
             },
             "Address of place was unexpectedly {val:?} for pointee type {ty:?}",
diff --git a/tests/codegen/mir-aggregate-no-alloca.rs b/tests/codegen/mir-aggregate-no-alloca.rs
index a7752d714fe..c0e7e1a05e3 100644
--- a/tests/codegen/mir-aggregate-no-alloca.rs
+++ b/tests/codegen/mir-aggregate-no-alloca.rs
@@ -55,7 +55,7 @@ pub fn make_cell_of_bool(b: bool) -> std::cell::Cell<bool> {
     std::cell::Cell::new(b)
 }
 
-// CHECK-LABLE: { i8, i16 } @make_cell_of_bool_and_short(i1 noundef zeroext %b, i16 noundef %s)
+// CHECK-LABEL: { i8, i16 } @make_cell_of_bool_and_short(i1 noundef zeroext %b, i16 noundef %s)
 #[no_mangle]
 pub fn make_cell_of_bool_and_short(b: bool, s: u16) -> std::cell::Cell<(bool, u16)> {
     // CHECK-NOT: alloca