Rollup merge of #69837 - jonas-schievink:gen-discr-opt, r=tmandry

Use smaller discriminants for generators

Closes https://github.com/rust-lang/rust/issues/69815

I'm not yet sure about the runtime performance impact of this, so I'll try running this on some benchmarks (if I can find any). (Update: No impact on the benchmarks I've measured on)

* [x] Add test with a generator that has exactly 256 total states
* [x] Add test with a generator that has more than 256 states so that it needs to use a u16 discriminant
* [x] Add tests for the size of `Option<[generator]>`
* [x] Add tests for the `discriminant_value` intrinsic in all cases

7 files changed, 167 insertions, 30 deletions

diff --git a/src/librustc/ty/layout.rs b/src/librustc/ty/layout.rs
index dedb3035ced..f5ef9bda0ee 100644
--- a/src/librustc/ty/layout.rs
+++ b/src/librustc/ty/layout.rs
@@ -1409,12 +1409,15 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
         // locals as part of the prefix. We compute the layout of all of
         // these fields at once to get optimal packing.
         let discr_index = substs.as_generator().prefix_tys(def_id, tcx).count();
-        // FIXME(eddyb) set the correct vaidity range for the discriminant.
-        let discr_layout = self.layout_of(substs.as_generator().discr_ty(tcx))?;
-        let discr = match &discr_layout.abi {
-            Abi::Scalar(s) => s.clone(),
-            _ => bug!(),
-        };
+
+        // `info.variant_fields` already accounts for the reserved variants, so no need to add them.
+        let max_discr = (info.variant_fields.len() - 1) as u128;
+        let discr_int = Integer::fit_unsigned(max_discr);
+        let discr_int_ty = discr_int.to_ty(tcx, false);
+        let discr = Scalar { value: Primitive::Int(discr_int, false), valid_range: 0..=max_discr };
+        let discr_layout = self.tcx.intern_layout(LayoutDetails::scalar(self, discr.clone()));
+        let discr_layout = TyLayout { ty: discr_int_ty, details: discr_layout };
+
         let promoted_layouts = ineligible_locals
             .iter()
             .map(|local| subst_field(info.field_tys[local]))
diff --git a/src/test/ui/async-await/async-fn-size-moved-locals.rs b/src/test/ui/async-await/async-fn-size-moved-locals.rs
index 4a413381aa3..636fafc2bc4 100644
--- a/src/test/ui/async-await/async-fn-size-moved-locals.rs
+++ b/src/test/ui/async-await/async-fn-size-moved-locals.rs
@@ -110,9 +110,9 @@ async fn mixed_sizes() {
 }
 
 fn main() {
-    assert_eq!(1028, std::mem::size_of_val(&single()));
-    assert_eq!(1032, std::mem::size_of_val(&single_with_noop()));
-    assert_eq!(3084, std::mem::size_of_val(&joined()));
-    assert_eq!(3084, std::mem::size_of_val(&joined_with_noop()));
-    assert_eq!(7188, std::mem::size_of_val(&mixed_sizes()));
+    assert_eq!(1025, std::mem::size_of_val(&single()));
+    assert_eq!(1026, std::mem::size_of_val(&single_with_noop()));
+    assert_eq!(3078, std::mem::size_of_val(&joined()));
+    assert_eq!(3079, std::mem::size_of_val(&joined_with_noop()));
+    assert_eq!(7181, std::mem::size_of_val(&mixed_sizes()));
 }
diff --git a/src/test/ui/async-await/async-fn-size-uninit-locals.rs b/src/test/ui/async-await/async-fn-size-uninit-locals.rs
index 0558084f4f8..d5d7b3fc3f0 100644
--- a/src/test/ui/async-await/async-fn-size-uninit-locals.rs
+++ b/src/test/ui/async-await/async-fn-size-uninit-locals.rs
@@ -95,9 +95,9 @@ async fn join_retval() -> Joiner {
 }
 
 fn main() {
-    assert_eq!(8, std::mem::size_of_val(&single()));
-    assert_eq!(12, std::mem::size_of_val(&single_with_noop()));
-    assert_eq!(3084, std::mem::size_of_val(&joined()));
-    assert_eq!(3084, std::mem::size_of_val(&joined_with_noop()));
-    assert_eq!(3080, std::mem::size_of_val(&join_retval()));
+    assert_eq!(2, std::mem::size_of_val(&single()));
+    assert_eq!(3, std::mem::size_of_val(&single_with_noop()));
+    assert_eq!(3078, std::mem::size_of_val(&joined()));
+    assert_eq!(3078, std::mem::size_of_val(&joined_with_noop()));
+    assert_eq!(3074, std::mem::size_of_val(&join_retval()));
 }
diff --git a/src/test/ui/async-await/async-fn-size.rs b/src/test/ui/async-await/async-fn-size.rs
index b313992db4e..0c1f3636446 100644
--- a/src/test/ui/async-await/async-fn-size.rs
+++ b/src/test/ui/async-await/async-fn-size.rs
@@ -86,13 +86,13 @@ async fn await3_level5() -> u8 {
 
 fn main() {
     assert_eq!(2, std::mem::size_of_val(&base()));
-    assert_eq!(8, std::mem::size_of_val(&await1_level1()));
-    assert_eq!(12, std::mem::size_of_val(&await2_level1()));
-    assert_eq!(12, std::mem::size_of_val(&await3_level1()));
-    assert_eq!(24, std::mem::size_of_val(&await3_level2()));
-    assert_eq!(36, std::mem::size_of_val(&await3_level3()));
-    assert_eq!(48, std::mem::size_of_val(&await3_level4()));
-    assert_eq!(60, std::mem::size_of_val(&await3_level5()));
+    assert_eq!(3, std::mem::size_of_val(&await1_level1()));
+    assert_eq!(4, std::mem::size_of_val(&await2_level1()));
+    assert_eq!(5, std::mem::size_of_val(&await3_level1()));
+    assert_eq!(8, std::mem::size_of_val(&await3_level2()));
+    assert_eq!(11, std::mem::size_of_val(&await3_level3()));
+    assert_eq!(14, std::mem::size_of_val(&await3_level4()));
+    assert_eq!(17, std::mem::size_of_val(&await3_level5()));
 
     assert_eq!(1,   wait(base()));
     assert_eq!(1,   wait(await1_level1()));
diff --git a/src/test/ui/generator/discriminant.rs b/src/test/ui/generator/discriminant.rs
new file mode 100644
index 00000000000..8a0f8a380ab
--- /dev/null
+++ b/src/test/ui/generator/discriminant.rs
@@ -0,0 +1,134 @@
+//! Tests that generator discriminant sizes and ranges are chosen optimally and that they are
+//! reflected in the output of `mem::discriminant`.
+
+// run-pass
+
+#![feature(generators, generator_trait, core_intrinsics)]
+
+use std::intrinsics::discriminant_value;
+use std::marker::Unpin;
+use std::mem::size_of_val;
+use std::{cmp, ops::*};
+
+macro_rules! yield25 {
+    ($e:expr) => {
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+        yield $e;
+    };
+}
+
+/// Yields 250 times.
+macro_rules! yield250 {
+    () => {
+        yield250!(())
+    };
+
+    ($e:expr) => {
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+        yield25!($e);
+    };
+}
+
+fn cycle(gen: impl Generator<()> + Unpin, expected_max_discr: u64) {
+    let mut gen = Box::pin(gen);
+    let mut max_discr = 0;
+    loop {
+        max_discr = cmp::max(max_discr, discriminant_value(gen.as_mut().get_mut()));
+        match gen.as_mut().resume(()) {
+            GeneratorState::Yielded(_) => {}
+            GeneratorState::Complete(_) => {
+                assert_eq!(max_discr, expected_max_discr);
+                return;
+            }
+        }
+    }
+}
+
+fn main() {
+    // Has only one invalid discr. value.
+    let gen_u8_tiny_niche = || {
+        || {
+            // 3 reserved variants
+
+            yield250!(); // 253 variants
+
+            yield; // 254
+            yield; // 255
+        }
+    };
+
+    // Uses all values in the u8 discriminant.
+    let gen_u8_full = || {
+        || {
+            // 3 reserved variants
+
+            yield250!(); // 253 variants
+
+            yield; // 254
+            yield; // 255
+            yield; // 256
+        }
+    };
+
+    // Barely needs a u16 discriminant.
+    let gen_u16 = || {
+        || {
+            // 3 reserved variants
+
+            yield250!(); // 253 variants
+
+            yield; // 254
+            yield; // 255
+            yield; // 256
+            yield; // 257
+        }
+    };
+
+    assert_eq!(size_of_val(&gen_u8_tiny_niche()), 1);
+    assert_eq!(size_of_val(&Some(gen_u8_tiny_niche())), 1); // uses niche
+    assert_eq!(size_of_val(&Some(Some(gen_u8_tiny_niche()))), 2); // cannot use niche anymore
+    assert_eq!(size_of_val(&gen_u8_full()), 1);
+    assert_eq!(size_of_val(&Some(gen_u8_full())), 2); // cannot use niche
+    assert_eq!(size_of_val(&gen_u16()), 2);
+    assert_eq!(size_of_val(&Some(gen_u16())), 2); // uses niche
+
+    cycle(gen_u8_tiny_niche(), 254);
+    cycle(gen_u8_full(), 255);
+    cycle(gen_u16(), 256);
+}
diff --git a/src/test/ui/generator/resume-arg-size.rs b/src/test/ui/generator/resume-arg-size.rs
index ffdc98d6f19..4f08ac0702b 100644
--- a/src/test/ui/generator/resume-arg-size.rs
+++ b/src/test/ui/generator/resume-arg-size.rs
@@ -23,6 +23,6 @@ fn main() {
 
     // Neither of these generators have the resume arg live across the `yield`, so they should be
     // 4 Bytes in size (only storing the discriminant)
-    assert_eq!(size_of_val(&gen_copy), 4);
-    assert_eq!(size_of_val(&gen_move), 4);
+    assert_eq!(size_of_val(&gen_copy), 1);
+    assert_eq!(size_of_val(&gen_move), 1);
 }
diff --git a/src/test/ui/generator/size-moved-locals.rs b/src/test/ui/generator/size-moved-locals.rs
index 2864fbb2f3c..74c60d98154 100644
--- a/src/test/ui/generator/size-moved-locals.rs
+++ b/src/test/ui/generator/size-moved-locals.rs
@@ -58,7 +58,7 @@ fn overlap_move_points() -> impl Generator<Yield = (), Return = ()> {
     }
 }
 
-fn overlap_x_and_y() -> impl Generator<Yield = (), Return = ()>{
+fn overlap_x_and_y() -> impl Generator<Yield = (), Return = ()> {
     static || {
         let x = Foo([0; FOO_SIZE]);
         yield;
@@ -70,8 +70,8 @@ fn overlap_x_and_y() -> impl Generator<Yield = (), Return = ()>{
 }
 
 fn main() {
-    assert_eq!(1028, std::mem::size_of_val(&move_before_yield()));
-    assert_eq!(1032, std::mem::size_of_val(&move_before_yield_with_noop()));
-    assert_eq!(2056, std::mem::size_of_val(&overlap_move_points()));
-    assert_eq!(1032, std::mem::size_of_val(&overlap_x_and_y()));
+    assert_eq!(1025, std::mem::size_of_val(&move_before_yield()));
+    assert_eq!(1026, std::mem::size_of_val(&move_before_yield_with_noop()));
+    assert_eq!(2051, std::mem::size_of_val(&overlap_move_points()));
+    assert_eq!(1026, std::mem::size_of_val(&overlap_x_and_y()));
 }