Also allow constructing enums as SSA values

scottmcm · scottmcm · commit 019c41a9a818 · 2025-03-16T16:51:33.000-07:00
This means that simple enum usages, like the ones internal to `for i in 0..n` stay entirely in SSA, no allocas needed.
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -779,16 +779,25 @@ 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 ty = rvalue.ty(self.mir, self.cx.tcx());
                 let ty = self.monomorphize(ty);
                 let layout = self.cx.layout_of(ty);
 
+                let variant_idx = match **kind {
+                    mir::AggregateKind::Adt(_, vi, _, _, None) => vi,
+                    mir::AggregateKind::RawPtr(..)
+                    | mir::AggregateKind::Closure(..)
+                    | mir::AggregateKind::Tuple => FIRST_VARIANT,
+                    _ => bug!("{rvalue:?} in codegen_rvalue_operand"),
+                };
+                let variant_layout = layout.for_variant(self.cx, variant_idx);
+
                 // `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 variant_layout.fields.index_by_increasing_offset() {
                     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| {
@@ -801,12 +810,66 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 }
 
                 let output_scalars = self.value_kind(layout).scalars().unwrap();
+                match layout.variants {
+                    abi::Variants::Empty => bug!(),
+                    abi::Variants::Single { index } => assert_eq!(index, variant_idx),
+                    abi::Variants::Multiple { tag, ref tag_encoding, tag_field, .. } => {
+                        match *tag_encoding {
+                            abi::TagEncoding::Direct => {
+                                assert_eq!(layout.fields.offset(tag_field), abi::Size::ZERO);
+
+                                let discr = ty
+                                    .discriminant_for_variant(self.cx.tcx(), variant_idx)
+                                    .unwrap();
+                                let discr_layout = bx.cx().layout_of(discr.ty);
+                                let discr_ibty = self.cx.immediate_backend_type(discr_layout);
+                                let discr_const = self.cx.const_uint_big(discr_ibty, discr.val);
+
+                                let tag_layout = layout.field(self.cx, tag_field);
+                                let tag_ibty = self.cx.immediate_backend_type(tag_layout);
+                                let tag_const = bx.intcast(discr_const, tag_ibty, tag.is_signed());
+
+                                inputs.insert(0, tag_const);
+                                input_scalars.insert(0, tag);
+
+                                // If this variant has fewer payloads, add the extra `undef`s.
+                                if input_scalars.len() != output_scalars.len() {
+                                    // Since we just inserted a tag, this can only happen
+                                    // if we're missing the second thing in a pair.
+                                    assert_eq!(input_scalars.len(), 1);
+                                    assert_eq!(output_scalars.len(), 2);
+
+                                    let ibty =
+                                        self.cx.scalar_pair_element_backend_type(layout, 1, true);
+                                    let undef = self.cx.const_undef(ibty);
+                                    inputs.push(undef);
+                                    input_scalars.push(output_scalars[1]);
+                                }
+                            }
+                            abi::TagEncoding::Niche { untagged_variant, .. } => {
+                                // For now we only handle variants where the niche is automatic,
+                                // but catch it if `rvalue_creates_operand` changes.
+                                assert_eq!(untagged_variant, variant_idx);
+                            }
+                        }
+                    }
+                }
+
                 itertools::izip!(&mut inputs, input_scalars, output_scalars).for_each(
                     |(v, in_s, out_s)| {
                         if in_s != out_s {
                             // We have to be really careful about bool here, because
                             // `(bool,)` stays i1 but `Cell<bool>` becomes i8.
                             *v = bx.from_immediate(*v);
+
+                            // Also we might be storing `usize` into a pointer field,
+                            // in which case we have to `without_provenance` it.
+                            if let abi::Primitive::Int(..) = in_s.primitive()
+                                && let abi::Primitive::Pointer(..) = out_s.primitive()
+                            {
+                                *v = bx.ptradd(bx.const_null(bx.type_ptr()), *v);
+                            }
+
                             *v = bx.to_immediate_scalar(*v, out_s);
                         }
                     },
@@ -1146,14 +1209,17 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             // (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 mut variant_idx = None;
                 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, ..) => {
+                    mir::AggregateKind::Adt(def_id, vidx, ..) => {
+                        variant_idx = Some(vidx);
                         let adt_def = self.cx.tcx().adt_def(def_id);
-                        adt_def.is_struct() && !adt_def.repr().simd()
+                        (adt_def.is_struct() || adt_def.is_enum())
+                            && !adt_def.repr().simd()
                     }
                     mir::AggregateKind::Closure(..) => true,
                     // FIXME: Can we do this for simple coroutines too?
@@ -1163,7 +1229,17 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     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)
+
+                    variant_idx.is_none_or(|v| !layout.for_variant(self.cx, v).uninhabited)
+                        && !self.cx.is_backend_ref(layout)
+                        && match &layout.variants {
+                        abi::Variants::Empty => bug!(),
+                        abi::Variants::Single { .. } => true,
+                        abi::Variants::Multiple { tag_encoding, ..} => match *tag_encoding {
+                            abi::TagEncoding::Direct => true,
+                            abi::TagEncoding::Niche { untagged_variant, ..} => Some(untagged_variant) == variant_idx,
+                        }
+                    }
                 }
             }
         }
diff --git a/tests/codegen/align-struct.rs b/tests/codegen/align-struct.rs
@@ -15,7 +15,7 @@ pub struct Nested64 {
 
 pub enum Enum4 {
     A(i32),
-    B(i32),
+    B(i32, i32),
 }
 
 pub enum Enum64 {
@@ -52,7 +52,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
 }
diff --git a/tests/codegen/common_prim_int_ptr.rs b/tests/codegen/common_prim_int_ptr.rs
@@ -11,7 +11,7 @@
 #[no_mangle]
 pub fn insert_int(x: usize) -> Result<usize, Box<()>> {
     // CHECK: start:
-    // CHECK-NEXT: inttoptr i{{[0-9]+}} %x to ptr
+    // CHECK-NEXT: getelementptr i8, ptr null, i{{[0-9]+}} %x
     // CHECK-NEXT: insertvalue
     // CHECK-NEXT: ret { i{{[0-9]+}}, ptr }
     Ok(x)
diff --git a/tests/codegen/enum/enum-aggregate.rs b/tests/codegen/enum/enum-aggregate.rs
@@ -0,0 +1,81 @@
+//@ compile-flags: -Copt-level=0 -Cno-prepopulate-passes -Cdebuginfo=0
+//@ min-llvm-version: 19
+//@ only-64bit
+
+#![crate_type = "lib"]
+
+use std::cmp::Ordering;
+use std::num::NonZero;
+use std::ptr::NonNull;
+
+#[no_mangle]
+fn make_some_bool(x: bool) -> Option<bool> {
+    // CHECK-LABEL: i8 @make_some_bool(i1 zeroext %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: %[[WIDER:.+]] = zext i1 %x to i8
+    // CHECK-NEXT: ret i8 %[[WIDER]]
+    Some(x)
+}
+
+#[no_mangle]
+fn make_none_bool() -> Option<bool> {
+    // FIXME: Aggregating a niched variant other than the untagged one isn't supported yet
+
+    // CHECK-LABEL: i8 @make_none_bool()
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: %_0 = alloca
+    // CHECK-NEXT: store i8 2, ptr %_0
+    None
+}
+
+#[no_mangle]
+fn make_some_ordering(x: Ordering) -> Option<Ordering> {
+    // CHECK-LABEL: i8 @make_some_ordering(i8 %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret i8 %x
+    Some(x)
+}
+
+#[no_mangle]
+fn make_some_u16(x: u16) -> Option<u16> {
+    // CHECK-LABEL: { i16, i16 } @make_some_u16(i16 %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: %0 = insertvalue { i16, i16 } { i16 1, i16 poison }, i16 %x, 1
+    // CHECK-NEXT: ret { i16, i16 } %0
+    Some(x)
+}
+
+#[no_mangle]
+fn make_none_u16() -> Option<u16> {
+    // CHECK-LABEL: { i16, i16 } @make_none_u16()
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret { i16, i16 } { i16 0, i16 undef }
+    None
+}
+
+#[no_mangle]
+fn make_some_nzu32(x: NonZero<u32>) -> Option<NonZero<u32>> {
+    // CHECK-LABEL: i32 @make_some_nzu32(i32 %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret i32 %x
+    Some(x)
+}
+
+#[no_mangle]
+fn make_ok_ptr(x: NonNull<u16>) -> Result<NonNull<u16>, usize> {
+    // CHECK-LABEL: { i64, ptr } @make_ok_ptr(ptr %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: %0 = insertvalue { i64, ptr } { i64 0, ptr poison }, ptr %x, 1
+    // CHECK-NEXT: ret { i64, ptr } %0
+    Ok(x)
+}
+
+#[no_mangle]
+fn make_ok_int(x: usize) -> Result<usize, NonNull<u16>> {
+    // CHECK-LABEL: { i64, ptr } @make_ok_int(i64 %x)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: %[[NOPROV:.+]] = getelementptr i8, ptr null, i64 %x
+    // CHECK-NEXT: %[[R:.+]] = insertvalue { i64, ptr } { i64 0, ptr poison }, ptr %[[NOPROV]], 1
+    // CHECK-NEXT: ret { i64, ptr } %[[R]]
+    Ok(x)
+}
diff --git a/tests/codegen/range-loop.rs b/tests/codegen/range-loop.rs
@@ -14,7 +14,6 @@ pub fn call_for_zero_to_n(n: u32, f: fn(u32)) {
     // CHECK: start:
     // CHECK-NOT: alloca
     // CHECK: %[[IND:.+]] = alloca [4 x i8]
-    // CHECK-NEXT: %[[ALWAYS_SOME_OPTION:.+]] = alloca
     // CHECK-NOT: alloca
     // CHECK: store i32 0, ptr %[[IND]],
     // CHECK: br label %[[HEAD:.+]]
@@ -31,10 +30,8 @@ pub fn call_for_zero_to_n(n: u32, f: fn(u32)) {
     // CHECK: %[[T2:.+]] = load i32, ptr %[[IND]],
     // CHECK: %[[T3:.+]] = add nuw i32 %[[T2]], 1
     // CHECK: store i32 %[[T3]], ptr %[[IND]],
-
-    // CHECK: store i32 %[[T2]]
-    // CHECK: %[[T4:.+]] = load i32
-    // CHECK: call void %f(i32{{.+}}%[[T4]])
+    // CHECK: call void %f(i32{{.+}}%[[T2]])
+    // CHECK: br label %[[HEAD]]
 
     for i in 0..n {
         f(i);
