Store generator field names in GeneratorLayout.

cjgillot · cjgillot · commit 3a1edd821238 · 2023-06-19T16:50:52.000Z
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
@@ -676,8 +676,7 @@ fn build_union_fields_for_direct_tag_generator<'ll, 'tcx>(
         _ => unreachable!(),
     };
 
-    let (generator_layout, state_specific_upvar_names) =
-        cx.tcx.generator_layout_and_saved_local_names(generator_def_id);
+    let generator_layout = cx.tcx.optimized_mir(generator_def_id).generator_layout().unwrap();
 
     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);
@@ -714,7 +713,6 @@ fn build_union_fields_for_direct_tag_generator<'ll, 'tcx>(
                 generator_type_and_layout,
                 generator_type_di_node,
                 generator_layout,
-                &state_specific_upvar_names,
                 &common_upvar_names,
             );
 
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs
@@ -6,7 +6,7 @@ use rustc_hir::def::CtorKind;
 use rustc_index::IndexSlice;
 use rustc_middle::{
     bug,
-    mir::{GeneratorLayout, GeneratorSavedLocal},
+    mir::GeneratorLayout,
     ty::{
         self,
         layout::{IntegerExt, LayoutOf, PrimitiveExt, TyAndLayout},
@@ -323,7 +323,6 @@ pub fn build_generator_variant_struct_type_di_node<'ll, 'tcx>(
     generator_type_and_layout: TyAndLayout<'tcx>,
     generator_type_di_node: &'ll DIType,
     generator_layout: &GeneratorLayout<'tcx>,
-    state_specific_upvar_names: &IndexSlice<GeneratorSavedLocal, Option<Symbol>>,
     common_upvar_names: &IndexSlice<FieldIdx, Symbol>,
 ) -> &'ll DIType {
     let variant_name = GeneratorSubsts::variant_name(variant_index);
@@ -357,7 +356,7 @@ pub fn build_generator_variant_struct_type_di_node<'ll, 'tcx>(
                 .map(|field_index| {
                     let generator_saved_local = generator_layout.variant_fields[variant_index]
                         [FieldIdx::from_usize(field_index)];
-                    let field_name_maybe = state_specific_upvar_names[generator_saved_local];
+                    let field_name_maybe = generator_layout.field_names[generator_saved_local];
                     let field_name = field_name_maybe
                         .as_ref()
                         .map(|s| Cow::from(s.as_str()))
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs
@@ -155,8 +155,8 @@ pub(super) fn build_generator_di_node<'ll, 'tcx>(
             DIFlags::FlagZero,
         ),
         |cx, generator_type_di_node| {
-            let (generator_layout, state_specific_upvar_names) =
-                cx.tcx.generator_layout_and_saved_local_names(generator_def_id);
+            let generator_layout =
+                cx.tcx.optimized_mir(generator_def_id).generator_layout().unwrap();
 
             let Variants::Multiple { tag_encoding: TagEncoding::Direct, ref variants, .. } = generator_type_and_layout.variants else {
                 bug!(
@@ -195,7 +195,6 @@ pub(super) fn build_generator_di_node<'ll, 'tcx>(
                                 generator_type_and_layout,
                                 generator_type_di_node,
                                 generator_layout,
-                                &state_specific_upvar_names,
                                 &common_upvar_names,
                             ),
                         source_info,
diff --git a/compiler/rustc_middle/src/mir/query.rs b/compiler/rustc_middle/src/mir/query.rs
@@ -9,6 +9,7 @@ use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
 use rustc_index::bit_set::BitMatrix;
 use rustc_index::{Idx, IndexVec};
+use rustc_span::symbol::Symbol;
 use rustc_span::Span;
 use rustc_target::abi::{FieldIdx, VariantIdx};
 use smallvec::SmallVec;
@@ -150,6 +151,9 @@ pub struct GeneratorLayout<'tcx> {
     /// The type of every local stored inside the generator.
     pub field_tys: IndexVec<GeneratorSavedLocal, GeneratorSavedTy<'tcx>>,
 
+    /// The name for debuginfo.
+    pub field_names: IndexVec<GeneratorSavedLocal, Option<Symbol>>,
+
     /// Which of the above fields are in each variant. Note that one field may
     /// be stored in multiple variants.
     pub variant_fields: IndexVec<VariantIdx, IndexVec<FieldIdx, GeneratorSavedLocal>>,
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
@@ -1,7 +1,6 @@
 //! 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::query::Providers;
 use crate::ty::layout::IntegerExt;
 use crate::ty::{
@@ -17,7 +16,6 @@ use rustc_hir as hir;
 use rustc_hir::def::{CtorOf, DefKind, Res};
 use rustc_hir::def_id::{CrateNum, DefId, LocalDefId};
 use rustc_index::bit_set::GrowableBitSet;
-use rustc_index::{Idx, IndexVec};
 use rustc_macros::HashStable;
 use rustc_session::Limit;
 use rustc_span::sym;
@@ -738,48 +736,6 @@ impl<'tcx> TyCtxt<'tcx> {
         if visitor.found_recursion { Err(expanded_type) } else { Ok(expanded_type) }
     }
 
-    // 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)
-    }
-
     /// Query and get an English description for the item's kind.
     pub fn def_descr(self, def_id: DefId) -> &'static str {
         self.def_kind_descr(self.def_kind(def_id), def_id)
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
@@ -230,7 +230,7 @@ struct TransformVisitor<'tcx> {
 
     // Mapping from Local to (type of local, generator struct index)
     // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
-    remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
+    remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, FieldIdx)>,
 
     // A map from a suspension point in a block to the locals which have live storage at that point
     storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
@@ -295,11 +295,11 @@ impl<'tcx> TransformVisitor<'tcx> {
     }
 
     // Create a Place referencing a generator struct field
-    fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> {
+    fn make_field(&self, variant_index: VariantIdx, idx: FieldIdx, ty: Ty<'tcx>) -> Place<'tcx> {
         let self_place = Place::from(SELF_ARG);
         let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index);
         let mut projection = base.projection.to_vec();
-        projection.push(ProjectionElem::Field(FieldIdx::new(idx), ty));
+        projection.push(ProjectionElem::Field(idx, ty));
 
         Place { local: base.local, projection: self.tcx.mk_place_elems(&projection) }
     }
@@ -904,7 +904,7 @@ fn compute_layout<'tcx>(
     liveness: LivenessInfo,
     body: &Body<'tcx>,
 ) -> (
-    FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
+    FxHashMap<Local, (Ty<'tcx>, VariantIdx, FieldIdx)>,
     GeneratorLayout<'tcx>,
     IndexVec<BasicBlock, Option<BitSet<Local>>>,
 ) {
@@ -982,6 +982,7 @@ fn compute_layout<'tcx>(
             // just use the first one here. That's fine; fields do not move
             // around inside generators, so it doesn't matter which variant
             // index we access them by.
+            let idx = FieldIdx::from_usize(idx);
             remap.entry(locals[saved_local]).or_insert((tys[saved_local].ty, variant_index, idx));
         }
         variant_fields.push(fields);
@@ -990,8 +991,23 @@ fn compute_layout<'tcx>(
     debug!("generator variant_fields = {:?}", variant_fields);
     debug!("generator storage_conflicts = {:#?}", storage_conflicts);
 
-    let layout =
-        GeneratorLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
+    let mut field_names = IndexVec::from_elem(None, &tys);
+    for var in &body.var_debug_info {
+        let VarDebugInfoContents::Place(place) = &var.value else { continue };
+        let Some(local) = place.as_local() else { continue };
+        let Some(&(_, variant, field)) = remap.get(&local) else { continue };
+
+        let saved_local = variant_fields[variant][field];
+        field_names.get_or_insert_with(saved_local, || var.name);
+    }
+
+    let layout = GeneratorLayout {
+        field_tys: tys,
+        field_names,
+        variant_fields,
+        variant_source_info,
+        storage_conflicts,
+    };
     debug!(?layout);
 
     (remap, layout, storage_liveness)
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
@@ -944,7 +944,7 @@ fn variant_info_for_generator<'tcx>(
         return (vec![], None);
     };
 
-    let (generator, state_specific_names) = cx.tcx.generator_layout_and_saved_local_names(def_id);
+    let generator = cx.tcx.optimized_mir(def_id).generator_layout().unwrap();
     let upvar_names = cx.tcx.closure_saved_names_of_captured_variables(def_id);
 
     let mut upvars_size = Size::ZERO;
@@ -983,9 +983,10 @@ fn variant_info_for_generator<'tcx>(
                     variant_size = variant_size.max(offset + field_layout.size);
                     FieldInfo {
                         kind: FieldKind::GeneratorLocal,
-                        name: state_specific_names.get(*local).copied().flatten().unwrap_or(
-                            Symbol::intern(&format!(".generator_field{}", local.as_usize())),
-                        ),
+                        name: generator.field_names[*local].unwrap_or(Symbol::intern(&format!(
+                            ".generator_field{}",
+                            local.as_usize()
+                        ))),
                         offset: offset.bytes(),
                         size: field_layout.size.bytes(),
                         align: field_layout.align.abi.bytes(),
