Use smaller discriminants for generators

Author: jonas-schievink

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]))

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()));
 }

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()));
 }

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()));

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()));
 }