Fix missed embedded-payload niche enum optimization (size enums)

compiler/rustc_abi/src/layout.rs

@@ -639,8 +639,22 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 last: all_indices.rev().find(|v| needs_disc(*v)).unwrap(),
             };
 
-            let count =
-                (niche_variants.last.index() as u128 - niche_variants.start.index() as u128) + 1;
+            let embedded_payload = {
+                let embedded = niche_variants.last;
+                let layout = &variant_layouts[embedded];
+                if let BackendRepr::Scalar(Scalar::Initialized { valid_range, .. }) =
+                    layout.backend_repr
+                    && layout.field_offsets.len() == 1
+                    && layout.field_offsets[FieldIdx::new(0)] == Size::ZERO
+                {
+                    Some((embedded, valid_range.end - valid_range.start + 1))
+                } else {
+                    None
+                }
+            };
+
+            let count = (niche_variants.last.index() as u128 - niche_variants.start.index() as u128)
+                + embedded_payload.map_or(1, |(_, states)| states);
 
             // Use the largest niche in the largest variant.
             let niche = variant_layouts[largest_variant_index].largest_niche?;
@@ -656,6 +670,11 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
                 layout.largest_niche = None;
 
+                if embedded_payload.is_some_and(|(embedded, _)| i == embedded) {
+                    // Its payload lives in the tag, so its own layout is unused.
+                    return true;
+                }
+
                 if layout.size <= niche_offset {
                     // This variant will fit before the niche.
                     return true;
@@ -731,6 +750,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                         untagged_variant: largest_variant_index,
                         niche_variants,
                         niche_start,
+                        embedded_payload,
                     },
                     tag_field: FieldIdx::new(0),
                     variants: variant_layouts,

compiler/rustc_abi/src/lib.rs

@@ -2030,6 +2030,10 @@ pub enum TagEncoding<VariantIdx: Idx> {
         /// This is inbounds of the type of the niche field
         /// (not sign-extended, i.e., all bits beyond the niche field size are 0).
         niche_start: u128,
+        /// The last `niche_variants` variant stores its payload in the tag:
+        /// `(variant, value_count)`. It must be a scalar with a primitive
+        /// field at offset 0. Its tag is `niche_start + (i - niche_variants.start) + payload`.
+        embedded_payload: Option<(VariantIdx, u128)>,
     },
 }
 

compiler/rustc_abi/src/tests.rs

@@ -68,3 +68,94 @@ fn wrapping_range_contains_range() {
     assert!(!boolr.contains_range(cmpr, size1));
     assert!(cmpr.contains_range(boolr, size1));
 }
+
+#[test]
+fn embedded_payload_niche_layout() {
+    let dl = TargetDataLayout::default();
+    let cx = LayoutCalculator::new(&dl);
+
+    let terminate = LayoutData::scalar(
+        &dl,
+        Scalar::Initialized {
+            value: Primitive::Int(Integer::I8, false),
+            valid_range: WrappingRange { start: 0, end: 1 },
+        },
+    );
+    let cleanup = LayoutData::scalar(
+        &dl,
+        Scalar::Initialized {
+            value: Primitive::Int(Integer::I32, false),
+            valid_range: WrappingRange { start: 0, end: 0xFFFF_FF00 },
+        },
+    );
+
+    #[derive(Copy, Clone)]
+    struct Field<'a>(&'a LayoutData<FieldIdx, VariantIdx>);
+
+    impl<'a> std::ops::Deref for Field<'a> {
+        type Target = &'a LayoutData<FieldIdx, VariantIdx>;
+        fn deref(&self) -> &Self::Target {
+            &self.0
+        }
+    }
+
+    impl<'a> std::fmt::Debug for Field<'a> {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            f.debug_struct("Field")
+                .field("size", &self.0.size.bytes())
+                .field("align", &self.0.align.abi.bytes())
+                .finish()
+        }
+    }
+
+    let unit: IndexVec<FieldIdx, Field<'_>> = IndexVec::new();
+    let variants = IndexVec::from_raw(vec![
+        unit.clone(),
+        unit,
+        IndexVec::from_raw(vec![Field(&terminate)]),
+        IndexVec::from_raw(vec![Field(&cleanup)]),
+    ]);
+
+    let layout = match cx.layout_of_struct_or_enum(
+        &ReprOptions::default(),
+        &variants,
+        true,
+        false,
+        |_, _| (Integer::I32, false),
+        (0..4).map(|i| (VariantIdx::new(i), i as u128)),
+        true,
+    ) {
+        Ok(layout) => layout,
+        Err(_) => panic!("layout calculation failed"),
+    };
+
+    assert_eq!(layout.size.bytes(), 4);
+    assert_eq!(layout.align, AbiAlign::new(Align::from_bytes(4).unwrap()));
+
+    match layout.variants {
+        Variants::Multiple { tag, tag_encoding, .. } => {
+            match tag {
+                Scalar::Initialized { value, valid_range } => {
+                    assert!(matches!(value, Primitive::Int(Integer::I32, false)));
+                    assert_eq!((valid_range.start, valid_range.end), (0, 0xFFFF_FF04 as u128));
+                }
+                _ => panic!("expected initialized tag"),
+            }
+            match tag_encoding {
+                TagEncoding::Niche {
+                    untagged_variant,
+                    niche_start,
+                    niche_variants,
+                    embedded_payload,
+                } => {
+                    assert_eq!(untagged_variant, VariantIdx::new(3));
+                    assert_eq!(niche_start, 0xFFFF_FF01);
+                    assert_eq!(niche_variants, (VariantIdx::new(0)..=VariantIdx::new(2)).into());
+                    assert_eq!(embedded_payload, Some((VariantIdx::new(2), 2)));
+                }
+                _ => panic!("expected niche encoding"),
+            }
+        }
+        _ => panic!("expected multiple variants"),
+    }
+}

compiler/rustc_codegen_cranelift/src/discriminant.rs

@@ -49,7 +49,7 @@ pub(crate) fn codegen_set_discriminant<'tcx>(
         Variants::Multiple {
             tag: _,
             tag_field,
-            tag_encoding: TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
+            tag_encoding: TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. },
             variants: _,
         } => {
             if variant_index != untagged_variant {
@@ -132,7 +132,7 @@ pub(crate) fn codegen_get_discriminant<'tcx>(
             let res = CValue::by_val(val, dest_layout);
             dest.write_cvalue(fx, res);
         }
-        TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start } => {
+        TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. } => {
             let relative_max = niche_variants.last.as_u32() - niche_variants.start.as_u32();
 
             // We have a subrange `niche_start..=niche_end` inside `range`.

compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs

@@ -406,7 +406,7 @@ fn compute_discriminant_value<'ll, 'tcx>(
             enum_type_and_layout.ty.discriminant_for_variant(cx.tcx, variant_index).unwrap().val,
         ),
         &Variants::Multiple {
-            tag_encoding: TagEncoding::Niche { ref niche_variants, niche_start, untagged_variant },
+            tag_encoding: TagEncoding::Niche { ref niche_variants, niche_start, untagged_variant, .. },
             tag,
             ..
         } => {

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -519,7 +519,7 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                 };
                 bx.intcast(tag_imm, cast_to, signed)
             }
-            TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start } => {
+            TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. } => {
                 // Cast to an integer so we don't have to treat a pointer as a
                 // special case.
                 let (tag, tag_llty) = match tag_scalar.primitive() {

compiler/rustc_codegen_ssa/src/mir/place.rs

@@ -497,7 +497,7 @@ pub(super) fn codegen_tag_value<'tcx, V>(
             Some((tag_field, imm))
         }
         Variants::Multiple {
-            tag_encoding: TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
+            tag_encoding: TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. },
             tag_field,
             ..
         } => {

compiler/rustc_const_eval/src/interpret/discriminant.rs

@@ -138,7 +138,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 // Return the cast value, and the index.
                 index.0
             }
-            TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start } => {
+            TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. } => {
                 let tag_val = tag_val.to_scalar();
                 // Compute the variant this niche value/"tag" corresponds to. With niche layout,
                 // discriminant (encoded in niche/tag) and variant index are the same.
@@ -298,7 +298,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
             abi::Variants::Multiple {
                 tag_encoding:
-                    TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
+                    TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start, .. },
                 tag: tag_layout,
                 tag_field,
                 ..

compiler/rustc_middle/src/ty/layout.rs

@@ -1116,7 +1116,7 @@ where
                     // dereferenceable.
                     Variants::Multiple {
                         tag_encoding:
-                            TagEncoding::Niche { untagged_variant, niche_variants, niche_start },
+                            TagEncoding::Niche { untagged_variant, niche_variants, niche_start, .. },
                         tag_field,
                         variants,
                         ..

compiler/rustc_public/src/abi.rs

@@ -229,6 +229,10 @@ pub enum TagEncoding {
         untagged_variant: VariantIdx,
         niche_variants: RangeInclusive<VariantIdx>,
         niche_start: u128,
+        /// If set, the last variant in `niche_variants` has an additional
+        /// scalar payload field at offset 0 that is stored in the niche
+        /// encoding (the "embedded payload" optimization).
+        embedded_payload: Option<(VariantIdx, u128)>,
     },
 }
 

compiler/rustc_public/src/unstable/convert/stable/abi.rs

@@ -238,11 +238,18 @@ impl<'tcx> Stable<'tcx> for rustc_abi::TagEncoding<rustc_abi::VariantIdx> {
     ) -> Self::T {
         match self {
             rustc_abi::TagEncoding::Direct => TagEncoding::Direct,
-            rustc_abi::TagEncoding::Niche { untagged_variant, niche_variants, niche_start } => {
+            rustc_abi::TagEncoding::Niche {
+                untagged_variant,
+                niche_variants,
+                niche_start,
+                embedded_payload,
+            } => {
                 TagEncoding::Niche {
                     untagged_variant: untagged_variant.stable(tables, cx),
                     niche_variants: niche_variants.stable(tables, cx),
                     niche_start: *niche_start,
+                    embedded_payload: embedded_payload
+                        .map(|(idx, payload)| (idx.stable(tables, cx), payload)),
                 }
             }
         }

compiler/rustc_ty_utils/src/layout/invariant.rs

@@ -297,7 +297,7 @@ pub(super) fn layout_sanity_check<'tcx>(cx: &LayoutCx<'tcx>, layout: &TyAndLayou
             }
         }
         Variants::Multiple { variants, tag, tag_encoding, .. } => {
-            if let TagEncoding::Niche { niche_start, untagged_variant, niche_variants } =
+            if let TagEncoding::Niche { niche_start, untagged_variant, niche_variants, .. } =
                 tag_encoding
             {
                 let niche_size = tag.size(cx);