Allow all MIR Aggregates to take the operand path (if layout permits)

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -571,6 +571,13 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
     pub(crate) fn builder(
         layout: TyAndLayout<'tcx>,
     ) -> Option<OperandRef<'tcx, Result<V, abi::Scalar>>> {
+        // Uninhabited types are weird, because for example `Result<!, !>`
+        // shows up as `FieldsShape::Primitive` and we need to be able to write
+        // a field into `(u32, !)`. We'll do that in an `alloca` instead.
+        if layout.uninhabited {
+            return None;
+        }
+
         let val = match layout.backend_repr {
             BackendRepr::Memory { .. } if layout.is_zst() => OperandValue::ZeroSized,
             BackendRepr::Scalar(s) => OperandValue::Immediate(Err(s)),
@@ -640,16 +647,46 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
         }
     }
 
+    /// Insert the immediate value `imm` for field `f` in the *type itself*,
+    /// rather than into one of the variants.
+    ///
+    /// Most things want [`OperandRef::insert_field`] instead, but this one is
+    /// necessary for writing things like enum tags that aren't in any variant.
+    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:?}"),
+        }
+    }
+
     /// After having set all necessary fields, this converts the
     /// `OperandValue<Result<V, _>>` (as obtained from [`OperandRef::builder`])
     /// to the normal `OperandValue<V>`.
     ///
     /// ICEs if any required fields were not set.
-    pub fn build(&self) -> OperandRef<'tcx, V> {
+    pub fn build(&self, cx: &impl CodegenMethods<'tcx, Value = V>) -> OperandRef<'tcx, V> {
         let OperandRef { val, layout } = *self;
 
+        // For something like `Option::<u32>::None`, it's expected that the
+        // payload scalar will not actually have been set, so this converts
+        // unset scalars to corresponding `undef` values so long as the scalar
+        // from the layout allows uninit.
         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::build called while fields are missing {self:?}"),
         };
 

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_tag_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_tag_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};
-use super::place::PlaceRef;
+use super::place::{PlaceRef, codegen_tag_value};
 use super::{FunctionCx, LocalRef};
 use crate::common::IntPredicate;
 use crate::traits::*;
@@ -694,7 +694,14 @@ 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);
@@ -706,10 +713,27 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 };
                 for (field_idx, field) in fields.iter_enumerated() {
                     let op = self.codegen_operand(bx, field);
-                    builder.insert_field(bx, FIRST_VARIANT, field_idx, op);
+                    let fi = active_field_index.unwrap_or(field_idx);
+                    builder.insert_field(bx, variant_index, fi, op);
                 }
 
-                builder.build()
+                let tag_result = codegen_tag_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.build(bx.cx())
+                    }
+                }
             }
             mir::Rvalue::ShallowInitBox(ref operand, content_ty) => {
                 let operand = self.codegen_operand(bx, operand);
@@ -1037,28 +1061,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             // Arrays are always aggregates, so it's not worth checking anything here.
             // (If it's really `[(); N]` or `[T; 0]` and we use the place path, fine.)
             mir::Rvalue::Repeat(..) => false,
-            mir::Rvalue::Aggregate(ref kind, _) => {
-                let allowed_kind = match **kind {
-                    // This always produces a `ty::RawPtr`, so will be Immediate or Pair
-                    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::Closure(..) => true,
-                    // FIXME: Can we do this for simple coroutines too?
-                    mir::AggregateKind::Coroutine(..) | mir::AggregateKind::CoroutineClosure(..) => false,
-                };
-                allowed_kind && {
+            mir::Rvalue::Aggregate(..) => {
                     let ty = rvalue.ty(self.mir, self.cx.tcx());
                     let ty = self.monomorphize(ty);
                     let layout = self.cx.spanned_layout_of(ty, span);
                     OperandRef::<Bx::Value>::builder(layout).is_some()
                 }
             }
-        }
 
         // (*) this is only true if the type is suitable
     }