Rollup merge of #104019 - compiler-errors:print-generator-sizes, r=wesleywiser

Compute generator sizes with `-Zprint_type_sizes`

Fixes #103887
r? `@pnkfelix`

6 files changed, 203 insertions, 104 deletions

diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
index d87117dffdc..a9e3dcf4cb3 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
@@ -27,9 +27,7 @@ use rustc_codegen_ssa::traits::*;
 use rustc_fs_util::path_to_c_string;
 use rustc_hir::def::CtorKind;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE};
-use rustc_index::vec::{Idx, IndexVec};
 use rustc_middle::bug;
-use rustc_middle::mir::{self, GeneratorLayout};
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::subst::GenericArgKind;
 use rustc_middle::ty::{
@@ -1026,33 +1024,6 @@ fn build_struct_type_di_node<'ll, 'tcx>(
 // Tuples
 //=-----------------------------------------------------------------------------
 
-/// Returns names of captured upvars for closures and generators.
-///
-/// Here are some examples:
-///  - `name__field1__field2` when the upvar is captured by value.
-///  - `_ref__name__field` when the upvar is captured by reference.
-///
-/// For generators this only contains upvars that are shared by all states.
-fn closure_saved_names_of_captured_variables(tcx: TyCtxt<'_>, def_id: DefId) -> SmallVec<String> {
-    let body = tcx.optimized_mir(def_id);
-
-    body.var_debug_info
-        .iter()
-        .filter_map(|var| {
-            let is_ref = match var.value {
-                mir::VarDebugInfoContents::Place(place) if place.local == mir::Local::new(1) => {
-                    // The projection is either `[.., Field, Deref]` or `[.., Field]`. It
-                    // implies whether the variable is captured by value or by reference.
-                    matches!(place.projection.last().unwrap(), mir::ProjectionElem::Deref)
-                }
-                _ => return None,
-            };
-            let prefix = if is_ref { "_ref__" } else { "" };
-            Some(prefix.to_owned() + var.name.as_str())
-        })
-        .collect()
-}
-
 /// Builds the DW_TAG_member debuginfo nodes for the upvars of a closure or generator.
 /// For a generator, this will handle upvars shared by all states.
 fn build_upvar_field_di_nodes<'ll, 'tcx>(
@@ -1083,7 +1054,7 @@ fn build_upvar_field_di_nodes<'ll, 'tcx>(
             .all(|&t| t == cx.tcx.normalize_erasing_regions(ParamEnv::reveal_all(), t))
     );
 
-    let capture_names = closure_saved_names_of_captured_variables(cx.tcx, def_id);
+    let capture_names = cx.tcx.closure_saved_names_of_captured_variables(def_id);
     let layout = cx.layout_of(closure_or_generator_ty);
 
     up_var_tys
@@ -1229,43 +1200,6 @@ fn build_union_type_di_node<'ll, 'tcx>(
     )
 }
 
-// FIXME(eddyb) maybe precompute this? Right now it's computed once
-// per generator monomorphization, but it doesn't depend on substs.
-fn generator_layout_and_saved_local_names<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    def_id: DefId,
-) -> (&'tcx GeneratorLayout<'tcx>, IndexVec<mir::GeneratorSavedLocal, Option<Symbol>>) {
-    let body = tcx.optimized_mir(def_id);
-    let generator_layout = body.generator_layout().unwrap();
-    let mut generator_saved_local_names = IndexVec::from_elem(None, &generator_layout.field_tys);
-
-    let state_arg = mir::Local::new(1);
-    for var in &body.var_debug_info {
-        let mir::VarDebugInfoContents::Place(place) = &var.value else { continue };
-        if place.local != state_arg {
-            continue;
-        }
-        match place.projection[..] {
-            [
-                // Deref of the `Pin<&mut Self>` state argument.
-                mir::ProjectionElem::Field(..),
-                mir::ProjectionElem::Deref,
-                // Field of a variant of the state.
-                mir::ProjectionElem::Downcast(_, variant),
-                mir::ProjectionElem::Field(field, _),
-            ] => {
-                let name = &mut generator_saved_local_names
-                    [generator_layout.variant_fields[variant][field]];
-                if name.is_none() {
-                    name.replace(var.name);
-                }
-            }
-            _ => {}
-        }
-    }
-    (generator_layout, generator_saved_local_names)
-}
-
 /// Computes the type parameters for a type, if any, for the given metadata.
 fn build_generic_type_param_di_nodes<'ll, 'tcx>(
     cx: &CodegenCx<'ll, 'tcx>,
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs
index 53e8a291d1e..69443b9b828 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs
@@ -22,9 +22,9 @@ use crate::{
     common::CodegenCx,
     debuginfo::{
         metadata::{
-            build_field_di_node, closure_saved_names_of_captured_variables,
+            build_field_di_node,
             enums::{tag_base_type, DiscrResult},
-            file_metadata, generator_layout_and_saved_local_names, size_and_align_of, type_di_node,
+            file_metadata, size_and_align_of, type_di_node,
             type_map::{self, Stub, UniqueTypeId},
             unknown_file_metadata, DINodeCreationResult, SmallVec, NO_GENERICS, NO_SCOPE_METADATA,
             UNKNOWN_LINE_NUMBER,
@@ -677,9 +677,9 @@ fn build_union_fields_for_direct_tag_generator<'ll, 'tcx>(
     };
 
     let (generator_layout, state_specific_upvar_names) =
-        generator_layout_and_saved_local_names(cx.tcx, generator_def_id);
+        cx.tcx.generator_layout_and_saved_local_names(generator_def_id);
 
-    let common_upvar_names = closure_saved_names_of_captured_variables(cx.tcx, generator_def_id);
+    let common_upvar_names = cx.tcx.closure_saved_names_of_captured_variables(generator_def_id);
     let variant_range = generator_substs.variant_range(generator_def_id, cx.tcx);
     let variant_count = (variant_range.start.as_u32()..variant_range.end.as_u32()).len();
 
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
index becbccc434d..93419d27a62 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
@@ -4,9 +4,8 @@ use crate::{
     common::CodegenCx,
     debuginfo::{
         metadata::{
-            closure_saved_names_of_captured_variables,
             enums::tag_base_type,
-            file_metadata, generator_layout_and_saved_local_names, size_and_align_of, type_di_node,
+            file_metadata, size_and_align_of, type_di_node,
             type_map::{self, Stub, StubInfo, UniqueTypeId},
             unknown_file_metadata, DINodeCreationResult, SmallVec, NO_GENERICS,
             UNKNOWN_LINE_NUMBER,
@@ -157,7 +156,7 @@ pub(super) fn build_generator_di_node<'ll, 'tcx>(
         ),
         |cx, generator_type_di_node| {
             let (generator_layout, state_specific_upvar_names) =
-                generator_layout_and_saved_local_names(cx.tcx, generator_def_id);
+                cx.tcx.generator_layout_and_saved_local_names(generator_def_id);
 
             let Variants::Multiple { tag_encoding: TagEncoding::Direct, ref variants, .. } = generator_type_and_layout.variants else {
                 bug!(
@@ -167,7 +166,7 @@ pub(super) fn build_generator_di_node<'ll, 'tcx>(
             };
 
             let common_upvar_names =
-                closure_saved_names_of_captured_variables(cx.tcx, generator_def_id);
+                cx.tcx.closure_saved_names_of_captured_variables(generator_def_id);
 
             // Build variant struct types
             let variant_struct_type_di_nodes: SmallVec<_> = variants
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
index 47c1ce80756..9ea8dc6e69f 100644
--- a/compiler/rustc_middle/src/ty/util.rs
+++ b/compiler/rustc_middle/src/ty/util.rs
@@ -1,6 +1,7 @@
 //! Miscellaneous type-system utilities that are too small to deserve their own modules.
 
 use crate::middle::codegen_fn_attrs::CodegenFnAttrFlags;
+use crate::mir;
 use crate::ty::layout::IntegerExt;
 use crate::ty::{
     self, DefIdTree, FallibleTypeFolder, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable,
@@ -15,6 +16,7 @@ use rustc_hir as hir;
 use rustc_hir::def::{CtorOf, DefKind, Res};
 use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::GrowableBitSet;
+use rustc_index::vec::{Idx, IndexVec};
 use rustc_macros::HashStable;
 use rustc_span::{sym, DUMMY_SP};
 use rustc_target::abi::{Integer, IntegerType, Size, TargetDataLayout};
@@ -692,6 +694,80 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn bound_impl_subject(self, def_id: DefId) -> ty::EarlyBinder<ty::ImplSubject<'tcx>> {
         ty::EarlyBinder(self.impl_subject(def_id))
     }
+
+    /// Returns names of captured upvars for closures and generators.
+    ///
+    /// Here are some examples:
+    ///  - `name__field1__field2` when the upvar is captured by value.
+    ///  - `_ref__name__field` when the upvar is captured by reference.
+    ///
+    /// For generators this only contains upvars that are shared by all states.
+    pub fn closure_saved_names_of_captured_variables(
+        self,
+        def_id: DefId,
+    ) -> SmallVec<[String; 16]> {
+        let body = self.optimized_mir(def_id);
+
+        body.var_debug_info
+            .iter()
+            .filter_map(|var| {
+                let is_ref = match var.value {
+                    mir::VarDebugInfoContents::Place(place)
+                        if place.local == mir::Local::new(1) =>
+                    {
+                        // The projection is either `[.., Field, Deref]` or `[.., Field]`. It
+                        // implies whether the variable is captured by value or by reference.
+                        matches!(place.projection.last().unwrap(), mir::ProjectionElem::Deref)
+                    }
+                    _ => return None,
+                };
+                let prefix = if is_ref { "_ref__" } else { "" };
+                Some(prefix.to_owned() + var.name.as_str())
+            })
+            .collect()
+    }
+
+    // FIXME(eddyb) maybe precompute this? Right now it's computed once
+    // per generator monomorphization, but it doesn't depend on substs.
+    pub fn generator_layout_and_saved_local_names(
+        self,
+        def_id: DefId,
+    ) -> (
+        &'tcx ty::GeneratorLayout<'tcx>,
+        IndexVec<mir::GeneratorSavedLocal, Option<rustc_span::Symbol>>,
+    ) {
+        let tcx = self;
+        let body = tcx.optimized_mir(def_id);
+        let generator_layout = body.generator_layout().unwrap();
+        let mut generator_saved_local_names =
+            IndexVec::from_elem(None, &generator_layout.field_tys);
+
+        let state_arg = mir::Local::new(1);
+        for var in &body.var_debug_info {
+            let mir::VarDebugInfoContents::Place(place) = &var.value else { continue };
+            if place.local != state_arg {
+                continue;
+            }
+            match place.projection[..] {
+                [
+                    // Deref of the `Pin<&mut Self>` state argument.
+                    mir::ProjectionElem::Field(..),
+                    mir::ProjectionElem::Deref,
+                    // Field of a variant of the state.
+                    mir::ProjectionElem::Downcast(_, variant),
+                    mir::ProjectionElem::Field(field, _),
+                ] => {
+                    let name = &mut generator_saved_local_names
+                        [generator_layout.variant_fields[variant][field]];
+                    if name.is_none() {
+                        name.replace(var.name);
+                    }
+                }
+                _ => {}
+            }
+        }
+        (generator_layout, generator_saved_local_names)
+    }
 }
 
 struct OpaqueTypeExpander<'tcx> {
diff --git a/compiler/rustc_session/src/code_stats.rs b/compiler/rustc_session/src/code_stats.rs
index eede4d16ea3..1085bce4475 100644
--- a/compiler/rustc_session/src/code_stats.rs
+++ b/compiler/rustc_session/src/code_stats.rs
@@ -19,7 +19,7 @@ pub enum SizeKind {
     Min,
 }
 
-#[derive(Clone, PartialEq, Eq, Hash, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
 pub struct FieldInfo {
     pub name: Symbol,
     pub offset: u64,
@@ -33,6 +33,7 @@ pub enum DataTypeKind {
     Union,
     Enum,
     Closure,
+    Generator,
 }
 
 #[derive(PartialEq, Eq, Hash, Debug)]
@@ -114,7 +115,7 @@ impl CodeStats {
 
             let struct_like = match kind {
                 DataTypeKind::Struct | DataTypeKind::Closure => true,
-                DataTypeKind::Enum | DataTypeKind::Union => false,
+                DataTypeKind::Enum | DataTypeKind::Union | DataTypeKind::Generator => false,
             };
             for (i, variant_info) in variants.into_iter().enumerate() {
                 let VariantInfo { ref name, kind: _, align: _, size, ref fields } = *variant_info;
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index fbc055b5d23..f4672a70072 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -1,3 +1,4 @@
+use hir::def_id::DefId;
 use rustc_hir as hir;
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::{Idx, IndexVec};
@@ -6,7 +7,7 @@ use rustc_middle::ty::layout::{
     IntegerExt, LayoutCx, LayoutError, LayoutOf, TyAndLayout, MAX_SIMD_LANES,
 };
 use rustc_middle::ty::{
-    self, subst::SubstsRef, EarlyBinder, ReprOptions, Ty, TyCtxt, TypeVisitable,
+    self, subst::SubstsRef, AdtDef, EarlyBinder, ReprOptions, Ty, TyCtxt, TypeVisitable,
 };
 use rustc_session::{DataTypeKind, FieldInfo, SizeKind, VariantInfo};
 use rustc_span::symbol::Symbol;
@@ -814,27 +815,39 @@ fn record_layout_for_printing_outlined<'tcx>(
         );
     };
 
-    let adt_def = match *layout.ty.kind() {
-        ty::Adt(ref adt_def, _) => {
+    match *layout.ty.kind() {
+        ty::Adt(adt_def, _) => {
             debug!("print-type-size t: `{:?}` process adt", layout.ty);
-            adt_def
+            let adt_kind = adt_def.adt_kind();
+            let adt_packed = adt_def.repr().pack.is_some();
+            let (variant_infos, opt_discr_size) = variant_info_for_adt(cx, layout, adt_def);
+            record(adt_kind.into(), adt_packed, opt_discr_size, variant_infos);
+        }
+
+        ty::Generator(def_id, substs, _) => {
+            debug!("print-type-size t: `{:?}` record generator", layout.ty);
+            // Generators always have a begin/poisoned/end state with additional suspend points
+            let (variant_infos, opt_discr_size) =
+                variant_info_for_generator(cx, layout, def_id, substs);
+            record(DataTypeKind::Generator, false, opt_discr_size, variant_infos);
         }
 
         ty::Closure(..) => {
             debug!("print-type-size t: `{:?}` record closure", layout.ty);
             record(DataTypeKind::Closure, false, None, vec![]);
-            return;
         }
 
         _ => {
             debug!("print-type-size t: `{:?}` skip non-nominal", layout.ty);
-            return;
         }
     };
+}
 
-    let adt_kind = adt_def.adt_kind();
-    let adt_packed = adt_def.repr().pack.is_some();
-
+fn variant_info_for_adt<'tcx>(
+    cx: &LayoutCx<'tcx, TyCtxt<'tcx>>,
+    layout: TyAndLayout<'tcx>,
+    adt_def: AdtDef<'tcx>,
+) -> (Vec<VariantInfo>, Option<Size>) {
     let build_variant_info = |n: Option<Symbol>, flds: &[Symbol], layout: TyAndLayout<'tcx>| {
         let mut min_size = Size::ZERO;
         let field_info: Vec<_> = flds
@@ -843,10 +856,7 @@ fn record_layout_for_printing_outlined<'tcx>(
             .map(|(i, &name)| {
                 let field_layout = layout.field(cx, i);
                 let offset = layout.fields.offset(i);
-                let field_end = offset + field_layout.size;
-                if min_size < field_end {
-                    min_size = field_end;
-                }
+                min_size = min_size.max(offset + field_layout.size);
                 FieldInfo {
                     name,
                     offset: offset.bytes(),
@@ -871,16 +881,9 @@ fn record_layout_for_printing_outlined<'tcx>(
                 debug!("print-type-size `{:#?}` variant {}", layout, adt_def.variant(index).name);
                 let variant_def = &adt_def.variant(index);
                 let fields: Vec<_> = variant_def.fields.iter().map(|f| f.name).collect();
-                record(
-                    adt_kind.into(),
-                    adt_packed,
-                    None,
-                    vec![build_variant_info(Some(variant_def.name), &fields, layout)],
-                );
+                (vec![build_variant_info(Some(variant_def.name), &fields, layout)], None)
             } else {
-                // (This case arises for *empty* enums; so give it
-                // zero variants.)
-                record(adt_kind.into(), adt_packed, None, vec![]);
+                (vec![], None)
             }
         }
 
@@ -898,15 +901,101 @@ fn record_layout_for_printing_outlined<'tcx>(
                     build_variant_info(Some(variant_def.name), &fields, layout.for_variant(cx, i))
                 })
                 .collect();
-            record(
-                adt_kind.into(),
-                adt_packed,
+
+            (
+                variant_infos,
                 match tag_encoding {
                     TagEncoding::Direct => Some(tag.size(cx)),
                     _ => None,
                 },
-                variant_infos,
-            );
+            )
         }
     }
 }
+
+fn variant_info_for_generator<'tcx>(
+    cx: &LayoutCx<'tcx, TyCtxt<'tcx>>,
+    layout: TyAndLayout<'tcx>,
+    def_id: DefId,
+    substs: ty::SubstsRef<'tcx>,
+) -> (Vec<VariantInfo>, Option<Size>) {
+    let Variants::Multiple { tag, ref tag_encoding, .. } = layout.variants else {
+        return (vec![], None);
+    };
+
+    let (generator, state_specific_names) = cx.tcx.generator_layout_and_saved_local_names(def_id);
+    let upvar_names = cx.tcx.closure_saved_names_of_captured_variables(def_id);
+
+    let mut upvars_size = Size::ZERO;
+    let upvar_fields: Vec<_> = substs
+        .as_generator()
+        .upvar_tys()
+        .zip(upvar_names)
+        .enumerate()
+        .map(|(field_idx, (_, name))| {
+            let field_layout = layout.field(cx, field_idx);
+            let offset = layout.fields.offset(field_idx);
+            upvars_size = upvars_size.max(offset + field_layout.size);
+            FieldInfo {
+                name: Symbol::intern(&name),
+                offset: offset.bytes(),
+                size: field_layout.size.bytes(),
+                align: field_layout.align.abi.bytes(),
+            }
+        })
+        .collect();
+
+    let variant_infos: Vec<_> = generator
+        .variant_fields
+        .iter_enumerated()
+        .map(|(variant_idx, variant_def)| {
+            let variant_layout = layout.for_variant(cx, variant_idx);
+            let mut variant_size = Size::ZERO;
+            let fields = variant_def
+                .iter()
+                .enumerate()
+                .map(|(field_idx, local)| {
+                    let field_layout = variant_layout.field(cx, field_idx);
+                    let offset = variant_layout.fields.offset(field_idx);
+                    // The struct is as large as the last field's end
+                    variant_size = variant_size.max(offset + field_layout.size);
+                    FieldInfo {
+                        name: state_specific_names.get(*local).copied().flatten().unwrap_or(
+                            Symbol::intern(&format!(".generator_field{}", local.as_usize())),
+                        ),
+                        offset: offset.bytes(),
+                        size: field_layout.size.bytes(),
+                        align: field_layout.align.abi.bytes(),
+                    }
+                })
+                .chain(upvar_fields.iter().copied())
+                .collect();
+
+            // If the variant has no state-specific fields, then it's the size of the upvars.
+            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,
+            };
+
+            VariantInfo {
+                name: Some(Symbol::intern(&ty::GeneratorSubsts::variant_name(variant_idx))),
+                kind: SizeKind::Exact,
+                size: variant_size.bytes(),
+                align: variant_layout.align.abi.bytes(),
+                fields,
+            }
+        })
+        .collect();
+    (
+        variant_infos,
+        match tag_encoding {
+            TagEncoding::Direct => Some(tag.size(cx)),
+            _ => None,
+        },
+    )
+}