Fix Async Generator ABI

This change was missed when making async generators implement `Future` directly.
It did not cause any problems in codegen so far, as `GeneratorState<(), Output>`
happens to have the same ABI as `Poll<Output>`.

1 files changed, 17 insertions, 5 deletions

diff --git a/compiler/rustc_ty_utils/src/abi.rs b/compiler/rustc_ty_utils/src/abi.rs
index 73c7eb6992f..d644cbccea1 100644
--- a/compiler/rustc_ty_utils/src/abi.rs
+++ b/compiler/rustc_ty_utils/src/abi.rs
@@ -85,7 +85,7 @@ fn fn_sig_for_fn_abi<'tcx>(
                 bound_vars,
             )
         }
-        ty::Generator(_, substs, _) => {
+        ty::Generator(did, substs, _) => {
             let sig = substs.as_generator().poly_sig();
 
             let bound_vars = tcx.mk_bound_variable_kinds(
@@ -104,10 +104,22 @@ fn fn_sig_for_fn_abi<'tcx>(
             let env_ty = tcx.mk_adt(pin_adt_ref, pin_substs);
 
             let sig = sig.skip_binder();
-            let state_did = tcx.require_lang_item(LangItem::GeneratorState, None);
-            let state_adt_ref = tcx.adt_def(state_did);
-            let state_substs = tcx.intern_substs(&[sig.yield_ty.into(), sig.return_ty.into()]);
-            let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
+            // The `FnSig` and the `ret_ty` here is for a generators main
+            // `Generator::resume(...) -> GeneratorState` function in case we
+            // have an ordinary generator, or the `Future::poll(...) -> Poll`
+            // function in case this is a special generator backing an async construct.
+            let ret_ty = if tcx.generator_is_async(did) {
+                let state_did = tcx.require_lang_item(LangItem::Poll, None);
+                let state_adt_ref = tcx.adt_def(state_did);
+                let state_substs = tcx.intern_substs(&[sig.return_ty.into()]);
+                tcx.mk_adt(state_adt_ref, state_substs)
+            } else {
+                let state_did = tcx.require_lang_item(LangItem::GeneratorState, None);
+                let state_adt_ref = tcx.adt_def(state_did);
+                let state_substs = tcx.intern_substs(&[sig.yield_ty.into(), sig.return_ty.into()]);
+                tcx.mk_adt(state_adt_ref, state_substs)
+            };
+
             ty::Binder::bind_with_vars(
                 tcx.mk_fn_sig(
                     [env_ty, sig.resume_ty].iter(),