Allow nearly all ADT aggregates to create operands

Author: scottmcm

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -2,7 +2,7 @@ use std::fmt;
 
 use rustc_abi as abi;
 use rustc_abi::{
-    Align, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, TagEncoding, Variants,
+    Align, BackendRepr, FIRST_VARIANT, FieldIdx, Primitive, Size, TagEncoding, VariantIdx, Variants,
 };
 use rustc_middle::mir::interpret::{Pointer, Scalar, alloc_range};
 use rustc_middle::mir::{self, ConstValue};
@@ -574,67 +574,109 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         };
         OperandRef { val, layout }
     }
+
+    pub(crate) fn supports_builder(layout: TyAndLayout<'tcx>) -> bool {
+        match layout.backend_repr {
+            BackendRepr::Memory { .. } if layout.is_zst() => true,
+            BackendRepr::Scalar(_) | BackendRepr::ScalarPair(_, _) => true,
+            _ => false,
+        }
+    }
 }
 
 impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
     pub(crate) fn insert_field<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         &mut self,
         bx: &mut Bx,
+        v: VariantIdx,
         f: FieldIdx,
         operand: OperandRef<'tcx, V>,
     ) {
-        let field_layout = self.layout.field(bx.cx(), f.as_usize());
-        let field_offset = self.layout.fields.offset(f.as_usize());
+        let (expect_zst, is_zero_offset) = if let abi::FieldsShape::Primitive = self.layout.fields {
+            // Don't ask for field layout for primitives, because that will panic.
+            if !self.layout.uninhabited {
+                // Real primitives only have one variant, but weird types like
+                // `Result<!, !>` turn out to also be "Primitive", and dead code
+                // like `Err(never)` needs to not ICE.
+                assert_eq!(v, FIRST_VARIANT);
+            }
+            let first_field = f == FieldIdx::ZERO;
+            (self.layout.is_zst() || !first_field, first_field)
+        } else {
+            let variant_layout = self.layout.for_variant(bx.cx(), v);
+            let field_layout = variant_layout.field(bx.cx(), f.as_usize());
+            let field_offset = variant_layout.fields.offset(f.as_usize());
+            (field_layout.is_zst(), field_offset == Size::ZERO)
+        };
 
         let mut update = |tgt: &mut Result<V, abi::Scalar>, src, from_scalar| {
             let from_bty = bx.cx().type_from_scalar(from_scalar);
             let to_scalar = tgt.unwrap_err();
             let to_bty = bx.cx().type_from_scalar(to_scalar);
-            let v = transmute_immediate(bx, src, from_scalar, from_bty, to_scalar, to_bty);
-            *tgt = Ok(v);
+            let imm = transmute_immediate(bx, src, from_scalar, from_bty, to_scalar, to_bty);
+            *tgt = Ok(imm);
         };
 
         match (operand.val, operand.layout.backend_repr) {
-            (OperandValue::ZeroSized, _) => {
-                debug_assert_eq!(field_layout.size, Size::ZERO);
-            }
+            (OperandValue::ZeroSized, _) if expect_zst => {}
             (OperandValue::Immediate(v), BackendRepr::Scalar(from_scalar)) => match &mut self.val {
-                OperandValue::Immediate(val @ Err(_)) => {
-                    debug_assert_eq!(field_offset, Size::ZERO);
+                OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
                     update(val, v, from_scalar);
-                    //*val = Ok(v);
                 }
-                OperandValue::Pair(fst @ Err(_), _) if field_offset == Size::ZERO => {
+                OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
                     update(fst, v, from_scalar);
-                    //*fst = Ok(v);
                 }
-                OperandValue::Pair(_, snd @ Err(_)) if field_offset != Size::ZERO => {
+                OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
                     update(snd, v, from_scalar);
-                    //*snd = Ok(v);
                 }
-                _ => bug!("Tried to insert {operand:?} into field {f:?} of {self:?}"),
+                _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
             },
             (OperandValue::Pair(a, b), BackendRepr::ScalarPair(from_sa, from_sb)) => {
                 match &mut self.val {
                     OperandValue::Pair(fst @ Err(_), snd @ Err(_)) => {
                         update(fst, a, from_sa);
-                        //*fst = Ok(a);
                         update(snd, b, from_sb);
-                        //*snd = Ok(b);
                     }
-                    _ => bug!("Tried to insert {operand:?} into field {f:?} of {self:?}"),
+                    _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
                 }
             }
-            _ => bug!("Unsupported operand {operand:?} inserting into field {f:?} of {self:?}"),
+            _ => bug!("Unsupported operand {operand:?} inserting into {v:?}.{f:?} of {self:?}"),
+        }
+    }
+
+    pub(super) fn insert_imm(&mut self, f: FieldIdx, imm: V) {
+        let field_offset = self.layout.fields.offset(f.as_usize());
+        let is_zero_offset = field_offset == Size::ZERO;
+        match &mut self.val {
+            OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
+                *val = Ok(imm);
+            }
+            OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
+                *fst = Ok(imm);
+            }
+            OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
+                *snd = Ok(imm);
+            }
+            _ => bug!("Tried to insert {imm:?} into field {f:?} of {self:?}"),
         }
     }
 
-    pub fn finalize(self) -> OperandRef<'tcx, V> {
-        let OperandRef { val, layout } = self;
+    pub fn finalize(&self, cx: &impl CodegenMethods<'tcx, Value = V>) -> OperandRef<'tcx, V> {
+        let OperandRef { val, layout } = *self;
+
+        let unwrap = |r: Result<V, abi::Scalar>| match r {
+            Ok(v) => v,
+            Err(s) if s.is_uninit_valid() => {
+                let bty = cx.type_from_scalar(s);
+                cx.const_undef(bty)
+            }
+            Err(_) => bug!("OperandRef::finalize called while fields are missing {self:?}"),
+        };
+
         let val = match val {
             OperandValue::ZeroSized => OperandValue::ZeroSized,
-            OperandValue::Immediate(v) => OperandValue::Immediate(v.unwrap()),
-            OperandValue::Pair(a, b) => OperandValue::Pair(a.unwrap(), b.unwrap()),
+            OperandValue::Immediate(v) => OperandValue::Immediate(unwrap(v)),
+            OperandValue::Pair(a, b) => OperandValue::Pair(unwrap(a), unwrap(b)),
             OperandValue::Ref(_) => bug!(),
         };
         OperandRef { val, layout }

compiler/rustc_codegen_ssa/src/mir/place.rs

@@ -1,4 +1,6 @@
-use rustc_abi::{Align, BackendRepr, FieldsShape, Size, TagEncoding, VariantIdx, Variants};
+use rustc_abi::{
+    Align, BackendRepr, FieldIdx, FieldsShape, Size, TagEncoding, VariantIdx, Variants,
+};
 use rustc_middle::mir::PlaceTy;
 use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
@@ -239,53 +241,17 @@ impl<'a, 'tcx, V: CodegenObject> PlaceRef<'tcx, V> {
         bx: &mut Bx,
         variant_index: VariantIdx,
     ) {
-        if self.layout.for_variant(bx.cx(), variant_index).is_uninhabited() {
-            // We play it safe by using a well-defined `abort`, but we could go for immediate UB
-            // if that turns out to be helpful.
-            bx.abort();
-            return;
-        }
-        match self.layout.variants {
-            Variants::Empty => unreachable!("we already handled uninhabited types"),
-            Variants::Single { index } => assert_eq!(index, variant_index),
-
-            Variants::Multiple { tag_encoding: TagEncoding::Direct, tag_field, .. } => {
-                let ptr = self.project_field(bx, tag_field.as_usize());
-                let to =
-                    self.layout.ty.discriminant_for_variant(bx.tcx(), variant_index).unwrap().val;
-                bx.store_to_place(
-                    bx.cx().const_uint_big(bx.cx().backend_type(ptr.layout), to),
-                    ptr.val,
-                );
+        match codegen_tagged_field_value(bx.cx(), variant_index, self.layout) {
+            Err(UninhabitedVariantError) => {
+                // We play it safe by using a well-defined `abort`, but we could go for immediate UB
+                // if that turns out to be helpful.
+                bx.abort();
             }
-            Variants::Multiple {
-                tag_encoding:
-                    TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
-                tag_field,
-                ..
-            } => {
-                if variant_index != untagged_variant {
-                    let niche = self.project_field(bx, tag_field.as_usize());
-                    let niche_llty = bx.cx().immediate_backend_type(niche.layout);
-                    let BackendRepr::Scalar(scalar) = niche.layout.backend_repr else {
-                        bug!("expected a scalar placeref for the niche");
-                    };
-                    // We are supposed to compute `niche_value.wrapping_add(niche_start)` wrapping
-                    // around the `niche`'s type.
-                    // The easiest way to do that is to do wrapping arithmetic on `u128` and then
-                    // masking off any extra bits that occur because we did the arithmetic with too many bits.
-                    let niche_value = variant_index.as_u32() - niche_variants.start().as_u32();
-                    let niche_value = (niche_value as u128).wrapping_add(niche_start);
-                    let niche_value = niche_value & niche.layout.size.unsigned_int_max();
-
-                    let niche_llval = bx.cx().scalar_to_backend(
-                        Scalar::from_uint(niche_value, niche.layout.size),
-                        scalar,
-                        niche_llty,
-                    );
-                    OperandValue::Immediate(niche_llval).store(bx, niche);
-                }
+            Ok(Some((tag_field, imm))) => {
+                let tag_place = self.project_field(bx, tag_field.as_usize());
+                OperandValue::Immediate(imm).store(bx, tag_place);
             }
+            Ok(None) => {}
         }
     }
 
@@ -471,3 +437,73 @@ fn round_up_const_value_to_alignment<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let offset = bx.and(neg_value, align_minus_1);
     bx.add(value, offset)
 }
+
+/// Calculates the value that needs to be stored to mark the discriminant.
+///
+/// This might be `None` for a `struct` or a niched variant (like `Some(&3)`).
+///
+/// If it's `Some`, it returns the value to store and the field in which to
+/// store it. Note that this value is *not* the same as the discriminant, in
+/// general, as it might be a niche value or have a different size.
+///
+/// It might also be an `Err` because the variant is uninhabited.
+pub(super) fn codegen_tagged_field_value<'tcx, V>(
+    cx: &impl CodegenMethods<'tcx, Value = V>,
+    variant_index: VariantIdx,
+    layout: TyAndLayout<'tcx>,
+) -> Result<Option<(FieldIdx, V)>, UninhabitedVariantError> {
+    // By checking uninhabited-ness first we don't need to worry about types
+    // like `(u32, !)` which are single-variant but weird.
+    if layout.for_variant(cx, variant_index).is_uninhabited() {
+        return Err(UninhabitedVariantError);
+    }
+
+    Ok(match layout.variants {
+        Variants::Empty => unreachable!("we already handled uninhabited types"),
+        Variants::Single { index } => {
+            assert_eq!(index, variant_index);
+            None
+        }
+
+        Variants::Multiple { tag_encoding: TagEncoding::Direct, tag_field, .. } => {
+            let discr = layout.ty.discriminant_for_variant(cx.tcx(), variant_index);
+            let to = discr.unwrap().val;
+            let tag_layout = layout.field(cx, tag_field.as_usize());
+            let tag_llty = cx.immediate_backend_type(tag_layout);
+            let imm = cx.const_uint_big(tag_llty, to);
+            Some((tag_field, imm))
+        }
+        Variants::Multiple {
+            tag_encoding: TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
+            tag_field,
+            ..
+        } => {
+            if variant_index != untagged_variant {
+                let niche_layout = layout.field(cx, tag_field.as_usize());
+                let niche_llty = cx.immediate_backend_type(niche_layout);
+                let BackendRepr::Scalar(scalar) = niche_layout.backend_repr else {
+                    bug!("expected a scalar placeref for the niche");
+                };
+                // We are supposed to compute `niche_value.wrapping_add(niche_start)` wrapping
+                // around the `niche`'s type.
+                // The easiest way to do that is to do wrapping arithmetic on `u128` and then
+                // masking off any extra bits that occur because we did the arithmetic with too many bits.
+                let niche_value = variant_index.as_u32() - niche_variants.start().as_u32();
+                let niche_value = (niche_value as u128).wrapping_add(niche_start);
+                let niche_value = niche_value & niche_layout.size.unsigned_int_max();
+
+                let niche_llval = cx.scalar_to_backend(
+                    Scalar::from_uint(niche_value, niche_layout.size),
+                    scalar,
+                    niche_llty,
+                );
+                Some((tag_field, niche_llval))
+            } else {
+                None
+            }
+        }
+    })
+}
+
+#[derive(Debug)]
+pub(super) struct UninhabitedVariantError;

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -10,7 +10,7 @@ use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument};
 
 use super::operand::{OperandRef, OperandValue, assume_scalar_range, transmute_immediate};
-use super::place::PlaceRef;
+use super::place::{PlaceRef, codegen_tagged_field_value};
 use super::{FunctionCx, LocalRef};
 use crate::common::IntPredicate;
 use crate::traits::*;
@@ -699,17 +699,42 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
             mir::Rvalue::Use(ref operand) => self.codegen_operand(bx, operand),
             mir::Rvalue::Repeat(..) => bug!("{rvalue:?} in codegen_rvalue_operand"),
-            mir::Rvalue::Aggregate(_, ref fields) => {
+            mir::Rvalue::Aggregate(ref kind, ref fields) => {
+                let (variant_index, active_field_index) = match **kind {
+                    mir::AggregateKind::Adt(_, variant_index, _, _, active_field_index) => {
+                        (variant_index, active_field_index)
+                    }
+                    _ => (FIRST_VARIANT, None),
+                };
+
                 let ty = rvalue.ty(self.mir, self.cx.tcx());
                 let ty = self.monomorphize(ty);
                 let layout = self.cx.layout_of(ty);
 
                 let mut builder = OperandRef::builder(layout);
                 for (field_idx, field) in fields.iter_enumerated() {
                     let op = self.codegen_operand(bx, field);
-                    builder.insert_field(bx, field_idx, op);
+                    let fi = active_field_index.unwrap_or(field_idx);
+                    builder.insert_field(bx, variant_index, fi, op);
+                }
+
+                let tag_result = codegen_tagged_field_value(self.cx, variant_index, layout);
+                match tag_result {
+                    Err(super::place::UninhabitedVariantError) => {
+                        // Like codegen_set_discr we use a sound abort, but could
+                        // potentially `unreachable` or just return the poison for
+                        // more optimizability, if that turns out to be helpful.
+                        bx.abort();
+                        let val = OperandValue::poison(bx, layout);
+                        OperandRef { val, layout }
+                    }
+                    Ok(maybe_tag_value) => {
+                        if let Some((tag_field, tag_imm)) = maybe_tag_value {
+                            builder.insert_imm(tag_field, tag_imm);
+                        }
+                        builder.finalize(bx.cx())
+                    }
                 }
-                builder.finalize()
             }
             mir::Rvalue::ShallowInitBox(ref operand, content_ty) => {
                 let operand = self.codegen_operand(bx, operand);
@@ -1043,19 +1068,16 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     mir::AggregateKind::RawPtr(..) => true,
                     mir::AggregateKind::Array(..) => false,
                     mir::AggregateKind::Tuple => true,
-                    mir::AggregateKind::Adt(def_id, ..) => {
-                        let adt_def = self.cx.tcx().adt_def(def_id);
-                        adt_def.is_struct() && !adt_def.repr().simd()
-                    }
+                    mir::AggregateKind::Adt(..) => true,
                     mir::AggregateKind::Closure(..) => true,
                     // FIXME: Can we do this for simple coroutines too?
                     mir::AggregateKind::Coroutine(..) | mir::AggregateKind::CoroutineClosure(..) => false,
                 };
                 allowed_kind && {
-                let ty = rvalue.ty(self.mir, self.cx.tcx());
-                let ty = self.monomorphize(ty);
+                    let ty = rvalue.ty(self.mir, self.cx.tcx());
+                    let ty = self.monomorphize(ty);
                     let layout = self.cx.spanned_layout_of(ty, span);
-                    !self.cx.is_backend_ref(layout)
+                    OperandRef::<Bx::Value>::supports_builder(layout)
                 }
             }
         }

tests/codegen/align-struct.rs

@@ -15,9 +15,11 @@ pub struct Nested64 {
     d: i8,
 }
 
+// This has the extra field in B to ensure it's not ScalarPair,
+// and thus that the test actually emits it via memory, not `insertvalue`.
 pub enum Enum4 {
     A(i32),
-    B(i32),
+    B(i32, i32),
 }
 
 pub enum Enum64 {
@@ -54,7 +56,7 @@ pub fn nested64(a: Align64, b: i32, c: i32, d: i8) -> Nested64 {
 // CHECK-LABEL: @enum4
 #[no_mangle]
 pub fn enum4(a: i32) -> Enum4 {
-    // CHECK: %e4 = alloca [8 x i8], align 4
+    // CHECK: %e4 = alloca [12 x i8], align 4
     let e4 = Enum4::A(a);
     e4
 }