Skip to content

Commit 2c22628

Browse files
azhoginazhogin
committed
Prereq6 for async drop - templated coroutine processing and layout
1 parent f9e319c commit 2c22628

File tree

12 files changed

+216
-41
lines changed

12 files changed

+216
-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
@@ -1310,7 +1310,11 @@ rustc_queries! {
13101310
/// Generates a MIR body for the shim.
13111311
query mir_shims(key: ty::InstanceKind<'tcx>) -> &'tcx mir::Body<'tcx> {
13121312
arena_cache
1313-
desc { |tcx| "generating MIR shim for `{}`", tcx.def_path_str(key.def_id()) }
1313+
desc {
1314+
|tcx| "generating MIR shim for `{}`, instance={:?}",
1315+
tcx.def_path_str(key.def_id()),
1316+
key
1317+
}
13141318
}
13151319

13161320
/// 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
@@ -922,23 +922,58 @@ where
922922
i,
923923
),
924924

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

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

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

+69-6
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ use rustc_hir::LangItem;
3737
use rustc_hir::def::{CtorKind, CtorOf, DefKind, DocLinkResMap, LifetimeRes, Res};
3838
use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, LocalDefIdMap};
3939
use rustc_index::IndexVec;
40+
use rustc_index::bit_set::BitMatrix;
4041
use rustc_macros::{
4142
Decodable, Encodable, HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable,
4243
extension,
@@ -100,7 +101,7 @@ pub use self::visit::{TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeV
100101
use crate::error::{OpaqueHiddenTypeMismatch, TypeMismatchReason};
101102
use crate::metadata::ModChild;
102103
use crate::middle::privacy::EffectiveVisibilities;
103-
use crate::mir::{Body, CoroutineLayout};
104+
use crate::mir::{Body, CoroutineLayout, CoroutineSavedLocal, CoroutineSavedTy, SourceInfo};
104105
use crate::query::{IntoQueryParam, Providers};
105106
use crate::ty;
106107
pub use crate::ty::diagnostics::*;
@@ -1740,9 +1741,8 @@ impl<'tcx> TyCtxt<'tcx> {
17401741
| ty::InstanceKind::CloneShim(..)
17411742
| ty::InstanceKind::ThreadLocalShim(..)
17421743
| ty::InstanceKind::FnPtrAddrShim(..)
1743-
| ty::InstanceKind::AsyncDropGlueCtorShim(..) => self.mir_shims(instance),
1744-
// async drop glue should be processed specifically, as a templated coroutine
1745-
ty::InstanceKind::AsyncDropGlue(_, _ty) => todo!(),
1744+
| ty::InstanceKind::AsyncDropGlueCtorShim(..)
1745+
| ty::InstanceKind::AsyncDropGlue(..) => self.mir_shims(instance),
17461746
}
17471747
}
17481748

@@ -1853,16 +1853,17 @@ impl<'tcx> TyCtxt<'tcx> {
18531853
self.def_kind(trait_def_id) == DefKind::TraitAlias
18541854
}
18551855

1856-
/// Returns layout of a coroutine. Layout might be unavailable if the
1856+
/// Returns layout of a non-templated coroutine. Layout might be unavailable if the
18571857
/// coroutine is tainted by errors.
18581858
///
18591859
/// Takes `coroutine_kind` which can be acquired from the `CoroutineArgs::kind_ty`,
18601860
/// e.g. `args.as_coroutine().kind_ty()`.
1861-
pub fn coroutine_layout(
1861+
pub fn ordinary_coroutine_layout(
18621862
self,
18631863
def_id: DefId,
18641864
coroutine_kind_ty: Ty<'tcx>,
18651865
) -> Option<&'tcx CoroutineLayout<'tcx>> {
1866+
debug_assert_ne!(Some(def_id), self.lang_items().async_drop_in_place_poll_fn());
18661867
let mir = self.optimized_mir(def_id);
18671868
// Regular coroutine
18681869
if coroutine_kind_ty.is_unit() {
@@ -1892,6 +1893,68 @@ impl<'tcx> TyCtxt<'tcx> {
18921893
}
18931894
}
18941895

1896+
/// Returns layout of a templated coroutine. Layout might be unavailable if the
1897+
/// coroutine is tainted by errors. Atm, the only templated coroutine is
1898+
/// `async_drop_in_place<T>::{closure}` returned from `async fn async_drop_in_place<T>(..)`.
1899+
pub fn templated_coroutine_layout(self, ty: Ty<'tcx>) -> Option<&'tcx CoroutineLayout<'tcx>> {
1900+
self.lang_items().async_drop_in_place_poll_fn().and_then(|def_id| {
1901+
self.mir_shims(InstanceKind::AsyncDropGlue(def_id, ty)).coroutine_layout_raw()
1902+
})
1903+
}
1904+
1905+
/// Returns layout of a templated (or not) coroutine. Layout might be unavailable if the
1906+
/// coroutine is tainted by errors.
1907+
pub fn coroutine_layout(
1908+
self,
1909+
def_id: DefId,
1910+
args: GenericArgsRef<'tcx>,
1911+
) -> Option<&'tcx CoroutineLayout<'tcx>> {
1912+
if Some(def_id) == self.lang_items().async_drop_in_place_poll_fn() {
1913+
fn find_impl_coroutine<'tcx>(tcx: TyCtxt<'tcx>, mut cor_ty: Ty<'tcx>) -> Ty<'tcx> {
1914+
let mut ty = cor_ty;
1915+
loop {
1916+
if let ty::Coroutine(def_id, args) = ty.kind() {
1917+
cor_ty = ty;
1918+
if tcx.is_templated_coroutine(*def_id) {
1919+
ty = args.first().unwrap().expect_ty();
1920+
continue;
1921+
} else {
1922+
return cor_ty;
1923+
}
1924+
} else {
1925+
return cor_ty;
1926+
}
1927+
}
1928+
}
1929+
// layout of `async_drop_in_place<T>::{closure}` in case,
1930+
// when T is a coroutine, contains this internal coroutine's ref
1931+
let arg_cor_ty = args.first().unwrap().expect_ty();
1932+
if arg_cor_ty.is_coroutine() {
1933+
let impl_cor_ty = find_impl_coroutine(self, arg_cor_ty);
1934+
let impl_ref = Ty::new_mut_ref(self, self.lifetimes.re_static, impl_cor_ty);
1935+
let span = self.def_span(def_id);
1936+
let source_info = SourceInfo::outermost(span);
1937+
let proxy_layout = CoroutineLayout {
1938+
field_tys: [CoroutineSavedTy {
1939+
ty: impl_ref,
1940+
source_info,
1941+
ignore_for_traits: true,
1942+
}]
1943+
.into(),
1944+
field_names: [None].into(),
1945+
variant_fields: [IndexVec::from([CoroutineSavedLocal::ZERO])].into(),
1946+
variant_source_info: [source_info].into(),
1947+
storage_conflicts: BitMatrix::new(1, 1),
1948+
};
1949+
return Some(self.arena.alloc(proxy_layout));
1950+
} else {
1951+
self.templated_coroutine_layout(Ty::new_coroutine(self, def_id, args))
1952+
}
1953+
} else {
1954+
self.ordinary_coroutine_layout(def_id, args.as_coroutine().kind_ty())
1955+
}
1956+
}
1957+
18951958
/// Given the `DefId` of an impl, returns the `DefId` of the trait it implements.
18961959
/// If it implements no trait, returns `None`.
18971960
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
@@ -78,8 +78,7 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
7878
#[inline]
7979
fn variant_range(&self, def_id: DefId, tcx: TyCtxt<'tcx>) -> Range<VariantIdx> {
8080
// FIXME requires optimized MIR
81-
FIRST_VARIANT
82-
..tcx.coroutine_layout(def_id, tcx.types.unit).unwrap().variant_fields.next_index()
81+
FIRST_VARIANT..tcx.coroutine_layout(def_id, self.args).unwrap().variant_fields.next_index()
8382
}
8483

8584
/// The discriminant for the given variant. Panics if the `variant_index` is
@@ -139,10 +138,14 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
139138
def_id: DefId,
140139
tcx: TyCtxt<'tcx>,
141140
) -> impl Iterator<Item: Iterator<Item = Ty<'tcx>> + Captures<'tcx>> {
142-
let layout = tcx.coroutine_layout(def_id, self.kind_ty()).unwrap();
141+
let layout = tcx.coroutine_layout(def_id, self.args).unwrap();
143142
layout.variant_fields.iter().map(move |variant| {
144143
variant.iter().map(move |field| {
145-
ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
144+
if tcx.is_templated_coroutine(def_id) {
145+
layout.field_tys[*field].ty
146+
} else {
147+
ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
148+
}
146149
})
147150
})
148151
}

‎compiler/rustc_mir_dataflow/src/value_analysis.rs

+3
Original file line numberDiff line numberDiff line change
@@ -407,6 +407,9 @@ impl<'tcx> Map<'tcx> {
407407
if exclude.contains(local) {
408408
continue;
409409
}
410+
if decl.ty.is_templated_coroutine(tcx) {
411+
continue;
412+
}
410413

411414
// Create a place for the local.
412415
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
@@ -624,7 +624,7 @@ fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
624624
}
625625
}
626626

627-
fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
627+
pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
628628
fn o1<T>(x: T) -> WithMinOptLevel<T> {
629629
WithMinOptLevel(1, x)
630630
}

‎compiler/rustc_mir_transform/src/validate.rs

+4-2
Original file line numberDiff line numberDiff line change
@@ -748,7 +748,9 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
748748
// since we may be in the process of computing this MIR in the
749749
// first place.
750750
let layout = if def_id == self.caller_body.source.def_id() {
751-
self.caller_body.coroutine_layout_raw()
751+
self.caller_body
752+
.coroutine_layout_raw()
753+
.or_else(|| self.tcx.coroutine_layout(def_id, args))
752754
} else if self.tcx.needs_coroutine_by_move_body_def_id(def_id)
753755
&& let ty::ClosureKind::FnOnce =
754756
args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap()
@@ -758,7 +760,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
758760
// Same if this is the by-move body of a coroutine-closure.
759761
self.caller_body.coroutine_layout_raw()
760762
} else {
761-
self.tcx.coroutine_layout(def_id, args.as_coroutine().kind_ty())
763+
self.tcx.coroutine_layout(def_id, args)
762764
};
763765

764766
let Some(layout) = layout else {

0 commit comments

Comments
 (0)