Make closure_saved_names_of_captured_variables a query.

cjgillot · cjgillot · commit 7d5b2e492638 · 2023-06-19T16:50:52.000Z
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/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs
@@ -324,7 +324,7 @@ pub fn build_generator_variant_struct_type_di_node<'ll, '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(
@@ -380,12 +380,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_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/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
@@ -738,38 +738,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(
diff --git a/compiler/rustc_mir_build/src/build/mod.rs b/compiler/rustc_mir_build/src/build/mod.rs
@@ -36,6 +36,29 @@ pub(crate) fn mir_built(
     tcx.alloc_steal_mir(mir_build(tcx, def))
 }
 
+/// 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(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_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
@@ -959,7 +959,7 @@ fn variant_info_for_generator<'tcx>(
             upvars_size = upvars_size.max(offset + field_layout.size);
             FieldInfo {
                 kind: FieldKind::Upvar,
-                name: Symbol::intern(&name),
+                name: *name,
                 offset: offset.bytes(),
                 size: field_layout.size.bytes(),
                 align: field_layout.align.abi.bytes(),

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,10 @@ impl<T> EraseType for &'_ ty::List<T> {`
`55`	`55`	`type Result = [u8; size_of::<*const ()>()];`
`56`	`56`	`}`
`57`	`57`
	`58`	`+impl<I: rustc_index::Idx, T> EraseType for &'_ rustc_index::IndexSlice<I, T> {`
	`59`	`+ type Result = [u8; size_of::<&'static rustc_index::IndexSlice<u32, ()>>()];`
	`60`	`+}`
	`61`	`+`
`58`	`62`	`impl<T> EraseType for Result<&'_ T, traits::query::NoSolution> {`
`59`	`63`	`type Result = [u8; size_of::<Result<&'static (), traits::query::NoSolution>>()];`
`60`	`64`	`}`