From 7f11d6f4bf8c4574b29d683ff72e596e7bcfbb50 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Wed, 29 May 2024 14:06:39 -0400
Subject: [PATCH] Add lang items for AsyncFn's associated types

---
 compiler/rustc_hir/src/lang_items.rs          |  3 +++
 compiler/rustc_span/src/symbol.rs             |  3 +++
 .../src/solve/normalizes_to/mod.rs            | 22 ++++++++++++-------
 library/core/src/ops/async_function.rs        |  3 +++
 4 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs
index c4be67cdd88..410c2ed8703 100644
--- a/compiler/rustc_hir/src/lang_items.rs
+++ b/compiler/rustc_hir/src/lang_items.rs
@@ -228,6 +228,9 @@ language_item_table! {
     AsyncFn,                 sym::async_fn,            async_fn_trait,             Target::Trait,          GenericRequirement::Exact(1);
     AsyncFnMut,              sym::async_fn_mut,        async_fn_mut_trait,         Target::Trait,          GenericRequirement::Exact(1);
     AsyncFnOnce,             sym::async_fn_once,       async_fn_once_trait,        Target::Trait,          GenericRequirement::Exact(1);
+    AsyncFnOnceOutput,       sym::async_fn_once_output,async_fn_once_output,       Target::AssocTy,        GenericRequirement::Exact(1);
+    CallOnceFuture,          sym::call_once_future,    call_once_future,           Target::AssocTy,        GenericRequirement::Exact(1);
+    CallRefFuture,           sym::call_ref_future,     call_ref_future,            Target::AssocTy,        GenericRequirement::Exact(2);
     AsyncFnKindHelper,       sym::async_fn_kind_helper,async_fn_kind_helper,       Target::Trait,          GenericRequirement::Exact(1);
 
     FnOnceOutput,            sym::fn_once_output,      fn_once_output,             Target::AssocTy,        GenericRequirement::None;
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 90da220b3f5..cabd462cc8a 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -441,6 +441,7 @@ symbols! {
         async_fn_kind_helper,
         async_fn_mut,
         async_fn_once,
+        async_fn_once_output,
         async_fn_track_caller,
         async_fn_traits,
         async_for_loop,
@@ -498,6 +499,8 @@ symbols! {
         call,
         call_mut,
         call_once,
+        call_once_future,
+        call_ref_future,
         caller_location,
         capture_disjoint_fields,
         catch_unwind,
diff --git a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
index 7fd2a3801cc..f9e164e1f3f 100644
--- a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
@@ -407,16 +407,20 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
                      output_coroutine_ty,
                      coroutine_return_ty,
                  }| {
-                    let (projection_term, term) = match tcx.item_name(goal.predicate.def_id()) {
-                        sym::CallOnceFuture => (
+                    let lang_items = tcx.lang_items();
+                    let (projection_term, term) = if Some(goal.predicate.def_id())
+                        == lang_items.call_once_future()
+                    {
+                        (
                             ty::AliasTerm::new(
                                 tcx,
                                 goal.predicate.def_id(),
                                 [goal.predicate.self_ty(), tupled_inputs_ty],
                             ),
                             output_coroutine_ty.into(),
-                        ),
-                        sym::CallRefFuture => (
+                        )
+                    } else if Some(goal.predicate.def_id()) == lang_items.call_ref_future() {
+                        (
                             ty::AliasTerm::new(
                                 tcx,
                                 goal.predicate.def_id(),
@@ -427,8 +431,9 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
                                 ],
                             ),
                             output_coroutine_ty.into(),
-                        ),
-                        sym::Output => (
+                        )
+                    } else if Some(goal.predicate.def_id()) == lang_items.async_fn_once_output() {
+                        (
                             ty::AliasTerm::new(
                                 tcx,
                                 goal.predicate.def_id(),
@@ -438,8 +443,9 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
                                 ],
                             ),
                             coroutine_return_ty.into(),
-                        ),
-                        name => bug!("no such associated type: {name}"),
+                        )
+                    } else {
+                        bug!("no such associated type in `AsyncFn*`: {:?}", goal.predicate.def_id())
                     };
                     ty::ProjectionPredicate { projection_term, term }
                 },
diff --git a/library/core/src/ops/async_function.rs b/library/core/src/ops/async_function.rs
index 18bcee5a1c7..fc21ffd5246 100644
--- a/library/core/src/ops/async_function.rs
+++ b/library/core/src/ops/async_function.rs
@@ -26,6 +26,7 @@ pub trait AsyncFn<Args: Tuple>: AsyncFnMut<Args> {
 pub trait AsyncFnMut<Args: Tuple>: AsyncFnOnce<Args> {
     /// Future returned by [`AsyncFnMut::async_call_mut`] and [`AsyncFn::async_call`].
     #[unstable(feature = "async_fn_traits", issue = "none")]
+    #[cfg_attr(not(bootstrap), lang = "call_ref_future")]
     type CallRefFuture<'a>: Future<Output = Self::Output>
     where
         Self: 'a;
@@ -46,10 +47,12 @@ pub trait AsyncFnMut<Args: Tuple>: AsyncFnOnce<Args> {
 pub trait AsyncFnOnce<Args: Tuple> {
     /// Future returned by [`AsyncFnOnce::async_call_once`].
     #[unstable(feature = "async_fn_traits", issue = "none")]
+    #[cfg_attr(not(bootstrap), lang = "call_once_future")]
     type CallOnceFuture: Future<Output = Self::Output>;
 
     /// Output type of the called closure's future.
     #[unstable(feature = "async_fn_traits", issue = "none")]
+    #[cfg_attr(not(bootstrap), lang = "async_fn_once_output")]
     type Output;
 
     /// Call the [`AsyncFnOnce`], returning a future which may move out of the called closure.