Rollup merge of #112759 - cjgillot:closure-names, r=oli-obk

Noratrieb · web-flow · commit 34c8e53d7a7f · 2023-06-21T07:37:01.000+02:00
Make closure_saved_names_of_captured_variables a query.

As we will start removing debuginfo during MIR optimizations, we need to keep them somewhere.
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs
@@ -1031,7 +1031,7 @@ fn build_upvar_field_di_nodes<'ll, 'tcx>(
             build_field_di_node(
                 cx,
                 closure_or_generator_di_node,
-                capture_name,
+                capture_name.as_str(),
                 cx.size_and_align_of(up_var_ty),
                 layout.fields.offset(index),
                 DIFlags::FlagZero,
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,8 +323,7 @@ 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: &[String],
+    common_upvar_names: &IndexSlice<FieldIdx, Symbol>,
 ) -> &'ll DIType {
     let variant_name = GeneratorSubsts::variant_name(variant_index);
     let unique_type_id = UniqueTypeId::for_enum_variant_struct_type(
@@ -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()))
@@ -380,12 +379,13 @@ pub fn build_generator_variant_struct_type_di_node<'ll, 'tcx>(
             // Fields that are common to all states
             let common_fields: SmallVec<_> = generator_substs
                 .prefix_tys()
+                .zip(common_upvar_names)
                 .enumerate()
-                .map(|(index, upvar_ty)| {
+                .map(|(index, (upvar_ty, upvar_name))| {
                     build_field_di_node(
                         cx,
                         variant_struct_type_di_node,
-                        &common_upvar_names[index],
+                        upvar_name.as_str(),
                         cx.size_and_align_of(upvar_ty),
                         generator_type_and_layout.fields.offset(index),
                         DIFlags::FlagZero,
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_data_structures/src/stable_hasher.rs b/compiler/rustc_data_structures/src/stable_hasher.rs
@@ -1,6 +1,6 @@
 use crate::sip128::SipHasher128;
 use rustc_index::bit_set::{self, BitSet};
-use rustc_index::{Idx, IndexVec};
+use rustc_index::{Idx, IndexSlice, IndexVec};
 use smallvec::SmallVec;
 use std::fmt;
 use std::hash::{BuildHasher, Hash, Hasher};
@@ -597,6 +597,18 @@ where
     }
 }
 
+impl<I: Idx, T, CTX> HashStable<CTX> for IndexSlice<I, T>
+where
+    T: HashStable<CTX>,
+{
+    fn hash_stable(&self, ctx: &mut CTX, hasher: &mut StableHasher) {
+        self.len().hash_stable(ctx, hasher);
+        for v in &self.raw {
+            v.hash_stable(ctx, hasher);
+        }
+    }
+}
+
 impl<I: Idx, T, CTX> HashStable<CTX> for IndexVec<I, T>
 where
     T: HashStable<CTX>,
diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -218,6 +218,7 @@ provide! { tcx, def_id, other, cdata,
     thir_abstract_const => { table }
     optimized_mir => { table }
     mir_for_ctfe => { table }
+    closure_saved_names_of_captured_variables => { table }
     mir_generator_witnesses => { table }
     promoted_mir => { table }
     def_span => { table }
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -1520,6 +1520,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             debug!("EntryBuilder::encode_mir({:?})", def_id);
             if encode_opt {
                 record!(self.tables.optimized_mir[def_id.to_def_id()] <- tcx.optimized_mir(def_id));
+                record!(self.tables.closure_saved_names_of_captured_variables[def_id.to_def_id()]
+                    <- tcx.closure_saved_names_of_captured_variables(def_id));
 
                 if tcx.sess.opts.unstable_opts.drop_tracking_mir
                     && let DefKind::Generator = self.tcx.def_kind(def_id)
diff --git a/compiler/rustc_metadata/src/rmeta/mod.rs b/compiler/rustc_metadata/src/rmeta/mod.rs
@@ -32,7 +32,7 @@ use rustc_span::edition::Edition;
 use rustc_span::hygiene::{ExpnIndex, MacroKind};
 use rustc_span::symbol::{Ident, Symbol};
 use rustc_span::{self, ExpnData, ExpnHash, ExpnId, Span};
-use rustc_target::abi::VariantIdx;
+use rustc_target::abi::{FieldIdx, VariantIdx};
 use rustc_target::spec::{PanicStrategy, TargetTriple};
 
 use std::marker::PhantomData;
@@ -416,6 +416,7 @@ define_tables! {
     object_lifetime_default: Table<DefIndex, LazyValue<ObjectLifetimeDefault>>,
     optimized_mir: Table<DefIndex, LazyValue<mir::Body<'static>>>,
     mir_for_ctfe: Table<DefIndex, LazyValue<mir::Body<'static>>>,
+    closure_saved_names_of_captured_variables: Table<DefIndex, LazyValue<IndexVec<FieldIdx, Symbol>>>,
     mir_generator_witnesses: Table<DefIndex, LazyValue<mir::GeneratorLayout<'static>>>,
     promoted_mir: Table<DefIndex, LazyValue<IndexVec<mir::Promoted, mir::Body<'static>>>>,
     thir_abstract_const: Table<DefIndex, LazyValue<ty::EarlyBinder<ty::Const<'static>>>>,
diff --git a/compiler/rustc_middle/src/arena.rs b/compiler/rustc_middle/src/arena.rs
@@ -27,6 +27,11 @@ macro_rules! arena_types {
                     rustc_middle::mir::Promoted,
                     rustc_middle::mir::Body<'tcx>
                 >,
+            [decode] closure_debuginfo:
+                rustc_index::IndexVec<
+                    rustc_target::abi::FieldIdx,
+                    rustc_span::symbol::Symbol,
+                >,
             [decode] typeck_results: rustc_middle::ty::TypeckResults<'tcx>,
             [decode] borrowck_result:
                 rustc_middle::mir::BorrowCheckResult<'tcx>,
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/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
@@ -55,6 +55,10 @@ impl<T> EraseType for &'_ ty::List<T> {
     type Result = [u8; size_of::<*const ()>()];
 }
 
+impl<I: rustc_index::Idx, T> EraseType for &'_ rustc_index::IndexSlice<I, T> {
+    type Result = [u8; size_of::<&'static rustc_index::IndexSlice<u32, ()>>()];
+}
+
 impl<T> EraseType for Result<&'_ T, traits::query::NoSolution> {
     type Result = [u8; size_of::<Result<&'static (), traits::query::NoSolution>>()];
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -531,6 +531,19 @@ rustc_queries! {
         }
     }
 
+    /// 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.
+    query closure_saved_names_of_captured_variables(def_id: DefId) -> &'tcx IndexVec<abi::FieldIdx, Symbol> {
+        arena_cache
+        desc { |tcx| "computing debuginfo for closure `{}`", tcx.def_path_str(def_id) }
+        separate_provide_extern
+    }
+
     query mir_generator_witnesses(key: DefId) -> &'tcx Option<mir::GeneratorLayout<'tcx>> {
         arena_cache
         desc { |tcx| "generator witness types for `{}`", tcx.def_path_str(key) }
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,80 +736,6 @@ impl<'tcx> TyCtxt<'tcx> {
         if visitor.found_recursion { Err(expanded_type) } else { Ok(expanded_type) }
     }
 
-    /// 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)
-    }
-
     /// 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_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
@@ -36,6 +36,22 @@ pub(crate) fn mir_built(
     tcx.alloc_steal_mir(mir_build(tcx, def))
 }
 
+pub(crate) fn closure_saved_names_of_captured_variables<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+) -> IndexVec<FieldIdx, Symbol> {
+    tcx.closure_captures(def_id)
+        .iter()
+        .map(|captured_place| {
+            let name = captured_place.to_symbol();
+            match captured_place.info.capture_kind {
+                ty::UpvarCapture::ByValue => name,
+                ty::UpvarCapture::ByRef(..) => Symbol::intern(&format!("_ref__{name}")),
+            }
+        })
+        .collect()
+}
+
 /// Construct the MIR for a given `DefId`.
 fn mir_build(tcx: TyCtxt<'_>, def: LocalDefId) -> Body<'_> {
     // Ensure unsafeck and abstract const building is ran before we steal the THIR.
diff --git a/compiler/rustc_mir_build/src/lib.rs b/compiler/rustc_mir_build/src/lib.rs
@@ -33,6 +33,8 @@ pub fn provide(providers: &mut Providers) {
     providers.check_match = thir::pattern::check_match;
     providers.lit_to_const = thir::constant::lit_to_const;
     providers.mir_built = build::mir_built;
+    providers.closure_saved_names_of_captured_variables =
+        build::closure_saved_names_of_captured_variables;
     providers.thir_check_unsafety = check_unsafety::thir_check_unsafety;
     providers.thir_body = thir::cx::thir_body;
     providers.thir_tree = thir::print::thir_tree;
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
