Fix enum layout to use niche-filling layout in guaranteed-NPO cases, even when the niche is in an uninhabited variant.

zachs18 · zachs18 · commit 018fe5a813fa · 2025-10-05T14:04:07.000-05:00
diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs
@@ -646,7 +646,10 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             return Err(LayoutCalculatorError::ReprConflict);
         }
 
-        let calculate_niche_filling_layout = || -> Option<LayoutData<FieldIdx, VariantIdx>> {
+        // Returns `(layout, is_this_maybe_npo)`.
+        // If `is_this_maybe_npo` is true, this layout is preferred over the tagged and
+        // no-tag layouts.
+        let calculate_niche_filling_layout = || -> Option<(LayoutData<_, _>, bool)> {
             if repr.inhibit_enum_layout_opt() {
                 return None;
             }
@@ -658,16 +661,12 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             let mut align = dl.aggregate_align;
             let mut max_repr_align = repr.align;
             let mut unadjusted_abi_align = align;
-            let mut inhabited_variants = 0;
 
             let mut variant_layouts = variants
                 .iter_enumerated()
                 .map(|(j, v)| {
                     let mut st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
                     st.variants = Variants::Single { index: j, variants: None };
-                    if !st.uninhabited {
-                        inhabited_variants += 1;
-                    }
 
                     align = align.max(st.align.abi);
                     max_repr_align = max_repr_align.max(st.max_repr_align);
@@ -677,41 +676,66 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 })
                 .collect::<Option<IndexVec<VariantIdx, _>>>()?;
 
-            if inhabited_variants < 2 {
-                // If there's only one inhabited variant, the no-tag layout will be equivalent to
-                // what the niched layout would be. Returning `None` here lets us assume there is
-                // another inhabited variant which simplifies the rest of the layout computation.
-                return None;
-            }
-
-            // Choose the largest variant, picking an inhabited variant in a tie
+            // Choose the largest variant, picking an inhabited variant in a tie.
+            // We still need to compute the niche-filling layout even if the largest variant
+            // is uninhabited, because `Result<UninhabitedReprTransparentPtr, 1ZST>` needs to
+            // still be NPO-optimized.
             let largest_variant_index = variant_layouts
                 .iter_enumerated()
                 .max_by_key(|(_i, layout)| (layout.size.bytes(), !layout.uninhabited))
                 .map(|(i, _layout)| i)?;
 
-            if variant_layouts[largest_variant_index].uninhabited {
-                // If the largest variant is uninhabited, then filling its niche would give
-                // a worse layout than the tagged layout.
-                return None;
-            }
+            // Use the largest niche in the largest variant.
+            let niche = variant_layouts[largest_variant_index].largest_niche?;
+            let niche_offset = niche.offset;
+            let niche_size = niche.value.size(dl);
+            let size = variant_layouts[largest_variant_index].size.align_to(align);
+            // If the niche occupies the whole size of the enum,
+            // this could be a case of NPO, so we should prefer this layout over the
+            // tagged and no-tag layouts, even if it has a worse niche (due to the
+            // niched variant being uninhabited)
+            let is_maybe_npo = niche_size == size;
 
             let all_indices = variants.indices();
             let needs_disc = |index: VariantIdx| {
                 index != largest_variant_index && !variant_layouts[index].uninhabited
             };
-            let niche_variants = all_indices.clone().find(|v| needs_disc(*v)).unwrap()
-                ..=all_indices.rev().find(|v| needs_disc(*v)).unwrap();
+            let first_niche_variant = all_indices.clone().find(|v| needs_disc(*v));
+            let last_niche_variant = all_indices.clone().rev().find(|v| needs_disc(*v));
+            let niche_variants = match Option::zip(first_niche_variant, last_niche_variant) {
+                Some((f, l)) => f..=l,
+                // All non-largest variants are uninhabited.
+                // If there are exactly 2 variants, and the largest variant's niche covers its
+                // whole size, and the non-largest variant is 1-aligned and zero-sized,
+                // this could be NPO.
+                // However, that only happens if the non-largest variant is "absent",
+                // which was already handled in `layout_of_struct_or_enum`, so we just debug_assert
+                // that that is not the case.
+                None if variants.len() == 2 && is_maybe_npo => {
+                    if cfg!(debug_assertions) {
+                        let non_largest_variant_index =
+                            all_indices.clone().find(|v| *v != largest_variant_index).unwrap();
+                        let non_largest_variant_layout =
+                            &variant_layouts[non_largest_variant_index];
+                        debug_assert!(
+                            non_largest_variant_layout.size > Size::ZERO
+                                || non_largest_variant_layout.align.abi > Align::ONE
+                        );
+                    }
+                    // The non-largest variant is not 1-ZST, this cannot be NPO,
+                    // use the no-tag layout.
+                    return None;
+                }
+                // All non-largest variants are uninhabited.
+                // This cannot be NPO, use the no-tag layout.
+                None => return None,
+            };
 
             let count =
                 (niche_variants.end().index() as u128 - niche_variants.start().index() as u128) + 1;
 
-            // Use the largest niche in the largest variant.
-            let niche = variant_layouts[largest_variant_index].largest_niche?;
+            // Caclulate the new niche.
             let (niche_start, niche_scalar) = niche.reserve(dl, count)?;
-            let niche_offset = niche.offset;
-            let niche_size = niche.value.size(dl);
-            let size = variant_layouts[largest_variant_index].size.align_to(align);
 
             let all_variants_fit = variant_layouts.iter_enumerated_mut().all(|(i, layout)| {
                 if i == largest_variant_index {
@@ -821,10 +845,19 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 randomization_seed: combined_seed,
             };
 
-            Some(layout)
+            Some((layout, is_maybe_npo))
         };
 
-        let niche_filling_layout = calculate_niche_filling_layout();
+        let niche_filling_layout = match calculate_niche_filling_layout() {
+            // If this is possibly NPO, we prefer this layout over the others,
+            // even if they might have a better niche.
+            // (They could never be smaller, since possibly-npo only occurs
+            // when the niche fills the whole size of the enum.)
+            Some((layout, true)) => return Ok(layout),
+            // This is definitely not NPO, try the other layouts.
+            Some((layout, false)) => Some(layout),
+            None => None,
+        };
 
         let discr_type = repr.discr_type();
         let discr_int = Integer::from_attr(dl, discr_type);
diff --git a/tests/ui/layout/enum.rs b/tests/ui/layout/enum.rs
@@ -115,3 +115,95 @@ enum AllUninhabitedVariantsAlignedZst { //~ERROR: size: Size(0 bytes)
     A(AlignedNever),
     B(AlignedNever),
 }
+
+
+#[repr(transparent)]
+struct NPONever(&'static (), !);
+
+#[repr(transparent)]
+struct NPONeverI16(std::num::NonZero<i16>, !);
+
+// All of these should be NPO-optimized, despite uninhabitedness of one or more variants
+#[rustc_layout(abi)]
+type NPONever1 = Result<NPONever, ()>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever2 = Result<(), NPONever>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever3 = Result<NPONever, !>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever4 = Result<!, NPONever>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever5 = Result<&'static (), !>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever6 = Result<!, &'static ()>;
+//~^ ERROR: abi: Scalar(Initialized { value: Pointer
+
+#[rustc_layout(abi)]
+type NPONever7 = Result<std::num::NonZero<i16>, !>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+type NPONever8 = Result<!, std::num::NonZero<i16>>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+type NPONever9 = Result<NPONeverI16, !>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+type NPONever10 = Result<!, NPONeverI16>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+type NPONever11 = Result<NPONeverI16, ()>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+type NPONever12 = Result<(), NPONeverI16>;
+//~^ERROR: abi: Scalar(Initialized { value: Int(I16
+
+#[rustc_layout(abi)]
+enum NPONever13 { //~ERROR: abi: Scalar(Initialized { value: Pointer
+    A(!, &'static (), !),
+    B((), !, [u8; 0]),
+}
+
+#[rustc_layout(abi)]
+enum NPONever14 { //~ERROR: abi: Scalar(Initialized { value: Pointer
+    A(!, &'static (), !),
+    B((), (), [u8; 0]),
+}
+
+#[rustc_layout(abi)]
+enum NPONever15 { //~ERROR: abi: Scalar(Initialized { value: Pointer
+    A((), &'static (), ()),
+    B((), !, [u8; 0]),
+}
+
+
+// These are not guaranteed to be NPO-optimized
+#[rustc_layout(abi)]
+type NotNPONever1 = Result<NPONever, NPONever>;
+//~^ERROR: abi: Scalar(
+
+#[rustc_layout(abi)]
+type NotNPONever2 = Result<NPONever, AlignedNever>;
+//~^ERROR: abi: Memory
+
+#[rustc_layout(abi)]
+type NotNPONever3 = Result<NPONever, &'static ()>;
+//~^ERROR: abi: Scalar(
+
+#[rustc_layout(abi)]
+type NotNPONever4 = Result<&'static (), AlignedNever>;
+//~^ERROR: abi: Scalar(
diff --git a/tests/ui/layout/enum.stderr b/tests/ui/layout/enum.stderr
@@ -242,5 +242,119 @@ error: abi: Memory { sized: true }
 LL | enum AllUninhabitedVariantsAlignedZst {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: aborting due to 22 previous errors
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: (..=0) | (1..) })
+  --> $DIR/enum.rs:128:1
+   |
+LL | type NPONever1 = Result<NPONever, ()>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: (..=0) | (1..) })
+  --> $DIR/enum.rs:132:1
+   |
+LL | type NPONever2 = Result<(), NPONever>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:136:1
+   |
+LL | type NPONever3 = Result<NPONever, !>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:140:1
+   |
+LL | type NPONever4 = Result<!, NPONever>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:144:1
+   |
+LL | type NPONever5 = Result<&'static (), !>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:148:1
+   |
+LL | type NPONever6 = Result<!, &'static ()>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: 1..=65535 })
+  --> $DIR/enum.rs:152:1
+   |
+LL | type NPONever7 = Result<std::num::NonZero<i16>, !>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: 1..=65535 })
+  --> $DIR/enum.rs:156:1
+   |
+LL | type NPONever8 = Result<!, std::num::NonZero<i16>>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: 1..=65535 })
+  --> $DIR/enum.rs:160:1
+   |
+LL | type NPONever9 = Result<NPONeverI16, !>;
+   | ^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: 1..=65535 })
+  --> $DIR/enum.rs:164:1
+   |
+LL | type NPONever10 = Result<!, NPONeverI16>;
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: (..=0) | (1..) })
+  --> $DIR/enum.rs:168:1
+   |
+LL | type NPONever11 = Result<NPONeverI16, ()>;
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I16, true), valid_range: (..=0) | (1..) })
+  --> $DIR/enum.rs:172:1
+   |
+LL | type NPONever12 = Result<(), NPONeverI16>;
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:176:1
+   |
+LL | enum NPONever13 {
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: (..=0) | (1..) })
+  --> $DIR/enum.rs:182:1
+   |
+LL | enum NPONever14 {
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:188:1
+   |
+LL | enum NPONever15 {
+   | ^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Int(I64, false), valid_range: 0..=0 })
+  --> $DIR/enum.rs:196:1
+   |
+LL | type NotNPONever1 = Result<NPONever, NPONever>;
+   | ^^^^^^^^^^^^^^^^^
+
+error: abi: Memory { sized: true }
+  --> $DIR/enum.rs:200:1
+   |
+LL | type NotNPONever2 = Result<NPONever, AlignedNever>;
+   | ^^^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:204:1
+   |
+LL | type NotNPONever3 = Result<NPONever, &'static ()>;
+   | ^^^^^^^^^^^^^^^^^
+
+error: abi: Scalar(Initialized { value: Pointer(AddressSpace(0)), valid_range: 1..=18446744073709551615 })
+  --> $DIR/enum.rs:208:1
+   |
+LL | type NotNPONever4 = Result<&'static (), AlignedNever>;
+   | ^^^^^^^^^^^^^^^^^
+
+error: aborting due to 41 previous errors
 
