Consider discriminant fields that are ordered before variant fields

Author: compiler-errors

compiler/rustc_ty_utils/src/layout.rs

@@ -919,7 +919,7 @@ fn variant_info_for_generator<'tcx>(
     def_id: DefId,
     substs: ty::SubstsRef<'tcx>,
 ) -> (Vec<VariantInfo>, Option<Size>) {
-    let Variants::Multiple { tag, ref tag_encoding, .. } = layout.variants else {
+    let Variants::Multiple { tag, ref tag_encoding, tag_field, .. } = layout.variants else {
         return (vec![], None);
     };
 
@@ -975,12 +975,28 @@ fn variant_info_for_generator<'tcx>(
             if variant_size == Size::ZERO {
                 variant_size = upvars_size;
             }
-            // We need to add the discriminant size back into min_size, since it is subtracted
-            // later during printing.
-            variant_size += match tag_encoding {
-                TagEncoding::Direct => tag.size(cx),
-                _ => Size::ZERO,
-            };
+
+            // This `if` deserves some explanation.
+            //
+            // The layout code has a choice of where to place the discriminant of this generator.
+            // If the discriminant of the generator is placed early in the layout (before the
+            // variant's own fields), then it'll implicitly be counted towards the size of the
+            // variant, since we use the maximum offset to calculate size.
+            //    (side-note: I know this is a bit problematic given upvars placement, etc).
+            //
+            // This is important, since the layout printing code always subtracts this discriminant
+            // size from the variant size if the struct is "enum"-like, so failing to account for it
+            // will either lead to numerical underflow, or an underreported variant size...
+            //
+            // However, if the discriminant is placed past the end of the variant, then we need
+            // to factor in the size of the discriminant manually. This really should be refactored
+            // better, but this "works" for now.
+            if layout.fields.offset(tag_field) >= variant_size {
+                variant_size += match tag_encoding {
+                    TagEncoding::Direct => tag.size(cx),
+                    _ => Size::ZERO,
+                };
+            }
 
             VariantInfo {
                 name: Some(Symbol::intern(&ty::GeneratorSubsts::variant_name(variant_idx))),

src/test/ui/print_type_sizes/generator_discr_placement.rs

@@ -0,0 +1,25 @@
+// compile-flags: -Z print-type-sizes
+// build-pass
+// ignore-pass
+
+// Tests a generator that has its discriminant as the *final* field.
+
+// Avoid emitting panic handlers, like the rest of these tests...
+#![feature(start, generators)]
+
+#[start]
+fn start(_: isize, _: *const *const u8) -> isize {
+    let a = || {
+        {
+            let w: i32 = 4;
+            yield;
+            drop(w);
+        }
+        {
+            let z: i32 = 7;
+            yield;
+            drop(z);
+        }
+    };
+    0
+}

src/test/ui/print_type_sizes/generator_discr_placement.stdout

@@ -0,0 +1,11 @@
+print-type-size type: `[generator@$DIR/generator_discr_placement.rs:12:13: 12:15]`: 8 bytes, alignment: 4 bytes
+print-type-size     discriminant: 1 bytes
+print-type-size     variant `Suspend0`: 7 bytes
+print-type-size         padding: 3 bytes
+print-type-size         field `.w`: 4 bytes, alignment: 4 bytes
+print-type-size     variant `Suspend1`: 7 bytes
+print-type-size         padding: 3 bytes
+print-type-size         field `.z`: 4 bytes, alignment: 4 bytes
+print-type-size     variant `Unresumed`: 0 bytes
+print-type-size     variant `Returned`: 0 bytes
+print-type-size     variant `Panicked`: 0 bytes