Move closure/generator type info methods to TyCtxt

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> {