Skip to content

Commit 08b025a

Browse files
azhoginazhogin
committed
Prereq6 for async drop - templated coroutine processing and layout
1 parent 9a0fb73 commit 08b025a

File tree

12 files changed

+208
-41
lines changed

12 files changed

+208
-41
lines changed

‎compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs

+1-2
Original file line numberDiff line numberDiff line change
@@ -721,8 +721,7 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
721721
_ => unreachable!(),
722722
};
723723

724-
let coroutine_layout =
725-
cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.kind_ty()).unwrap();
724+
let coroutine_layout = cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.args).unwrap();
726725

727726
let common_upvar_names = cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
728727
let variant_range = coroutine_args.variant_range(coroutine_def_id, cx.tcx);

‎compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs

+2-4
Original file line numberDiff line numberDiff line change
@@ -174,10 +174,8 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
174174
DIFlags::FlagZero,
175175
),
176176
|cx, coroutine_type_di_node| {
177-
let coroutine_layout = cx
178-
.tcx
179-
.coroutine_layout(coroutine_def_id, coroutine_args.as_coroutine().kind_ty())
180-
.unwrap();
177+
let coroutine_layout =
178+
cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args).unwrap();
181179

182180
let Variants::Multiple { tag_encoding: TagEncoding::Direct, ref variants, .. } =
183181
coroutine_type_and_layout.variants

‎compiler/rustc_middle/src/arena.rs

+1
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ macro_rules! arena_types {
99
($macro:path) => (
1010
$macro!([
1111
[] layout: rustc_abi::LayoutData<rustc_abi::FieldIdx, rustc_abi::VariantIdx>,
12+
[] proxy_coroutine_layout: rustc_middle::mir::CoroutineLayout<'tcx>,
1213
[] fn_abi: rustc_target::callconv::FnAbi<'tcx, rustc_middle::ty::Ty<'tcx>>,
1314
// AdtDef are interned and compared by address
1415
[decode] adt_def: rustc_middle::ty::AdtDefData,

‎compiler/rustc_middle/src/query/mod.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1344,7 +1344,11 @@ rustc_queries! {
13441344
/// Generates a MIR body for the shim.
13451345
query mir_shims(key: ty::InstanceKind<'tcx>) -> &'tcx mir::Body<'tcx> {
13461346
arena_cache
1347-
desc { |tcx| "generating MIR shim for `{}`", tcx.def_path_str(key.def_id()) }
1347+
desc {
1348+
|tcx| "generating MIR shim for `{}`, instance={:?}",
1349+
tcx.def_path_str(key.def_id()),
1350+
key
1351+
}
13481352
}
13491353

13501354
/// The `symbol_name` query provides the symbol name for calling a

‎compiler/rustc_middle/src/ty/layout.rs

+50-15
Original file line numberDiff line numberDiff line change
@@ -923,23 +923,58 @@ where
923923
i,
924924
),
925925

926-
ty::Coroutine(def_id, args) => match this.variants {
927-
Variants::Empty => unreachable!(),
928-
Variants::Single { index } => TyMaybeWithLayout::Ty(
929-
args.as_coroutine()
930-
.state_tys(def_id, tcx)
931-
.nth(index.as_usize())
932-
.unwrap()
933-
.nth(i)
934-
.unwrap(),
935-
),
936-
Variants::Multiple { tag, tag_field, .. } => {
937-
if i == tag_field {
938-
return TyMaybeWithLayout::TyAndLayout(tag_layout(tag));
926+
ty::Coroutine(def_id, args) => {
927+
// layout of `async_drop_in_place<T>::{closure}` in case,
928+
// when T is a coroutine, contains this internal coroutine's ref
929+
if tcx.is_templated_coroutine(def_id) {
930+
fn find_impl_coroutine<'tcx>(
931+
tcx: TyCtxt<'tcx>,
932+
mut cor_ty: Ty<'tcx>,
933+
) -> Ty<'tcx> {
934+
let mut ty = cor_ty;
935+
loop {
936+
if let ty::Coroutine(def_id, args) = ty.kind() {
937+
cor_ty = ty;
938+
if tcx.is_templated_coroutine(*def_id) {
939+
ty = args.first().unwrap().expect_ty();
940+
continue;
941+
} else {
942+
return cor_ty;
943+
}
944+
} else {
945+
return cor_ty;
946+
}
947+
}
948+
}
949+
let arg_cor_ty = args.first().unwrap().expect_ty();
950+
if arg_cor_ty.is_coroutine() {
951+
assert!(i == 0);
952+
let impl_cor_ty = find_impl_coroutine(tcx, arg_cor_ty);
953+
return TyMaybeWithLayout::Ty(Ty::new_mut_ref(
954+
tcx,
955+
tcx.lifetimes.re_static,
956+
impl_cor_ty,
957+
));
939958
}
940-
TyMaybeWithLayout::Ty(args.as_coroutine().prefix_tys()[i])
941959
}
942-
},
960+
match this.variants {
961+
Variants::Empty => unreachable!(),
962+
Variants::Single { index } => TyMaybeWithLayout::Ty(
963+
args.as_coroutine()
964+
.state_tys(def_id, tcx)
965+
.nth(index.as_usize())
966+
.unwrap()
967+
.nth(i)
968+
.unwrap(),
969+
),
970+
Variants::Multiple { tag, tag_field, .. } => {
971+
if i == tag_field {
972+
return TyMaybeWithLayout::TyAndLayout(tag_layout(tag));
973+
}
974+
TyMaybeWithLayout::Ty(args.as_coroutine().prefix_tys()[i])
975+
}
976+
}
977+
}
943978

944979
ty::Tuple(tys) => TyMaybeWithLayout::Ty(tys[i]),
945980

‎compiler/rustc_middle/src/ty/mod.rs

+69-6
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use rustc_hir::LangItem;
3939
use rustc_hir::def::{CtorKind, CtorOf, DefKind, DocLinkResMap, LifetimeRes, Res};
4040
use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, LocalDefIdMap};
4141
use rustc_index::IndexVec;
42+
use rustc_index::bit_set::BitMatrix;
4243
use rustc_macros::{
4344
Decodable, Encodable, HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable,
4445
extension,
@@ -103,7 +104,7 @@ pub use self::visit::*;
103104
use crate::error::{OpaqueHiddenTypeMismatch, TypeMismatchReason};
104105
use crate::metadata::ModChild;
105106
use crate::middle::privacy::EffectiveVisibilities;
106-
use crate::mir::{Body, CoroutineLayout};
107+
use crate::mir::{Body, CoroutineLayout, CoroutineSavedLocal, CoroutineSavedTy, SourceInfo};
107108
use crate::query::{IntoQueryParam, Providers};
108109
use crate::ty;
109110
use crate::ty::codec::{TyDecoder, TyEncoder};
@@ -1721,9 +1722,8 @@ impl<'tcx> TyCtxt<'tcx> {
17211722
| ty::InstanceKind::CloneShim(..)
17221723
| ty::InstanceKind::ThreadLocalShim(..)
17231724
| ty::InstanceKind::FnPtrAddrShim(..)
1724-
| ty::InstanceKind::AsyncDropGlueCtorShim(..) => self.mir_shims(instance),
1725-
// async drop glue should be processed specifically, as a templated coroutine
1726-
ty::InstanceKind::AsyncDropGlue(_, _ty) => todo!(),
1725+
| ty::InstanceKind::AsyncDropGlueCtorShim(..)
1726+
| ty::InstanceKind::AsyncDropGlue(..) => self.mir_shims(instance),
17271727
}
17281728
}
17291729

@@ -1839,16 +1839,17 @@ impl<'tcx> TyCtxt<'tcx> {
18391839
self.def_kind(trait_def_id) == DefKind::TraitAlias
18401840
}
18411841

1842-
/// Returns layout of a coroutine. Layout might be unavailable if the
1842+
/// Returns layout of a non-templated coroutine. Layout might be unavailable if the
18431843
/// coroutine is tainted by errors.
18441844
///
18451845
/// Takes `coroutine_kind` which can be acquired from the `CoroutineArgs::kind_ty`,
18461846
/// e.g. `args.as_coroutine().kind_ty()`.
1847-
pub fn coroutine_layout(
1847+
pub fn ordinary_coroutine_layout(
18481848
self,
18491849
def_id: DefId,
18501850
coroutine_kind_ty: Ty<'tcx>,
18511851
) -> Option<&'tcx CoroutineLayout<'tcx>> {
1852+
debug_assert_ne!(Some(def_id), self.lang_items().async_drop_in_place_poll_fn());
18521853
let mir = self.optimized_mir(def_id);
18531854
// Regular coroutine
18541855
if coroutine_kind_ty.is_unit() {
@@ -1878,6 +1879,68 @@ impl<'tcx> TyCtxt<'tcx> {
18781879
}
18791880
}
18801881

1882+
/// Returns layout of a templated coroutine. Layout might be unavailable if the
1883+
/// coroutine is tainted by errors. Atm, the only templated coroutine is
1884+
/// `async_drop_in_place<T>::{closure}` returned from `async fn async_drop_in_place<T>(..)`.
1885+
pub fn templated_coroutine_layout(self, ty: Ty<'tcx>) -> Option<&'tcx CoroutineLayout<'tcx>> {
1886+
self.lang_items().async_drop_in_place_poll_fn().and_then(|def_id| {
1887+
self.mir_shims(InstanceKind::AsyncDropGlue(def_id, ty)).coroutine_layout_raw()
1888+
})
1889+
}
1890+
1891+
/// Returns layout of a templated (or not) coroutine. Layout might be unavailable if the
1892+
/// coroutine is tainted by errors.
1893+
pub fn coroutine_layout(
1894+
self,
1895+
def_id: DefId,
1896+
args: GenericArgsRef<'tcx>,
1897+
) -> Option<&'tcx CoroutineLayout<'tcx>> {
1898+
if Some(def_id) == self.lang_items().async_drop_in_place_poll_fn() {
1899+
fn find_impl_coroutine<'tcx>(tcx: TyCtxt<'tcx>, mut cor_ty: Ty<'tcx>) -> Ty<'tcx> {
1900+
let mut ty = cor_ty;
1901+
loop {
1902+
if let ty::Coroutine(def_id, args) = ty.kind() {
1903+
cor_ty = ty;
1904+
if tcx.is_templated_coroutine(*def_id) {
1905+
ty = args.first().unwrap().expect_ty();
1906+
continue;
1907+
} else {
1908+
return cor_ty;
1909+
}
1910+
} else {
1911+
return cor_ty;
1912+
}
1913+
}
1914+
}
1915+
// layout of `async_drop_in_place<T>::{closure}` in case,
1916+
// when T is a coroutine, contains this internal coroutine's ref
1917+
let arg_cor_ty = args.first().unwrap().expect_ty();
1918+
if arg_cor_ty.is_coroutine() {
1919+
let impl_cor_ty = find_impl_coroutine(self, arg_cor_ty);
1920+
let impl_ref = Ty::new_mut_ref(self, self.lifetimes.re_static, impl_cor_ty);
1921+
let span = self.def_span(def_id);
1922+
let source_info = SourceInfo::outermost(span);
1923+
let proxy_layout = CoroutineLayout {
1924+
field_tys: [CoroutineSavedTy {
1925+
ty: impl_ref,
1926+
source_info,
1927+
ignore_for_traits: true,
1928+
}]
1929+
.into(),
1930+
field_names: [None].into(),
1931+
variant_fields: [IndexVec::from([CoroutineSavedLocal::ZERO])].into(),
1932+
variant_source_info: [source_info].into(),
1933+
storage_conflicts: BitMatrix::new(1, 1),
1934+
};
1935+
return Some(self.arena.alloc(proxy_layout));
1936+
} else {
1937+
self.templated_coroutine_layout(Ty::new_coroutine(self, def_id, args))
1938+
}
1939+
} else {
1940+
self.ordinary_coroutine_layout(def_id, args.as_coroutine().kind_ty())
1941+
}
1942+
}
1943+
18811944
/// Given the `DefId` of an impl, returns the `DefId` of the trait it implements.
18821945
/// If it implements no trait, returns `None`.
18831946
pub fn trait_id_of_impl(self, def_id: DefId) -> Option<DefId> {

‎compiler/rustc_middle/src/ty/sty.rs

+7-4
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,7 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
7676
#[inline]
7777
fn variant_range(&self, def_id: DefId, tcx: TyCtxt<'tcx>) -> Range<VariantIdx> {
7878
// FIXME requires optimized MIR
79-
FIRST_VARIANT
80-
..tcx.coroutine_layout(def_id, tcx.types.unit).unwrap().variant_fields.next_index()
79+
FIRST_VARIANT..tcx.coroutine_layout(def_id, self.args).unwrap().variant_fields.next_index()
8180
}
8281

8382
/// The discriminant for the given variant. Panics if the `variant_index` is
@@ -137,10 +136,14 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
137136
def_id: DefId,
138137
tcx: TyCtxt<'tcx>,
139138
) -> impl Iterator<Item: Iterator<Item = Ty<'tcx>>> {
140-
let layout = tcx.coroutine_layout(def_id, self.kind_ty()).unwrap();
139+
let layout = tcx.coroutine_layout(def_id, self.args).unwrap();
141140
layout.variant_fields.iter().map(move |variant| {
142141
variant.iter().map(move |field| {
143-
ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
142+
if tcx.is_templated_coroutine(def_id) {
143+
layout.field_tys[*field].ty
144+
} else {
145+
ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
146+
}
144147
})
145148
})
146149
}

‎compiler/rustc_mir_dataflow/src/value_analysis.rs

+3
Original file line numberDiff line numberDiff line change
@@ -405,6 +405,9 @@ impl<'tcx> Map<'tcx> {
405405
if exclude.contains(local) {
406406
continue;
407407
}
408+
if decl.ty.is_templated_coroutine(tcx) {
409+
continue;
410+
}
408411

409412
// Create a place for the local.
410413
debug_assert!(self.locals[local].is_none());

‎compiler/rustc_mir_transform/src/known_panics_lint.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -888,7 +888,11 @@ impl CanConstProp {
888888
};
889889
for (local, val) in cpv.can_const_prop.iter_enumerated_mut() {
890890
let ty = body.local_decls[local].ty;
891-
if ty.is_union() {
891+
if ty.is_templated_coroutine(tcx) {
892+
// No const propagation for templated coroutine (AsyncDropGlue)
893+
*val = ConstPropMode::NoPropagation;
894+
continue;
895+
} else if ty.is_union() {
892896
// Unions are incompatible with the current implementation of
893897
// const prop because Rust has no concept of an active
894898
// variant of a union

‎compiler/rustc_mir_transform/src/lib.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -648,7 +648,7 @@ fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
648648
}
649649
}
650650

651-
fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
651+
pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
652652
fn o1<T>(x: T) -> WithMinOptLevel<T> {
653653
WithMinOptLevel(1, x)
654654
}

‎compiler/rustc_mir_transform/src/validate.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -751,7 +751,9 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
751751
// since we may be in the process of computing this MIR in the
752752
// first place.
753753
let layout = if def_id == self.caller_body.source.def_id() {
754-
self.caller_body.coroutine_layout_raw()
754+
self.caller_body
755+
.coroutine_layout_raw()
756+
.or_else(|| self.tcx.coroutine_layout(def_id, args))
755757
} else if self.tcx.needs_coroutine_by_move_body_def_id(def_id)
756758
&& let ty::ClosureKind::FnOnce =
757759
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap()
@@ -761,7 +763,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
761763
// Same if this is the by-move body of a coroutine-closure.
762764
self.caller_body.coroutine_layout_raw()
763765
} else {
764-
self.tcx.coroutine_layout(def_id, args.as_coroutine().kind_ty())
766+
self.tcx.coroutine_layout(def_id, args)
765767
};
766768

767769
let Some(layout) = layout else {

0 commit comments

Comments
 (0)