Implement async gen blocks

compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

@@ -112,6 +112,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     self.assemble_future_candidates(obligation, &mut candidates);
                 } else if lang_items.iterator_trait() == Some(def_id) {
                     self.assemble_iterator_candidates(obligation, &mut candidates);
+                } else if lang_items.async_iterator_trait() == Some(def_id) {
+                    self.assemble_async_iterator_candidates(obligation, &mut candidates);
                 }
 
                 self.assemble_closure_candidates(obligation, &mut candidates);
@@ -258,6 +260,34 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         }
     }
 
+    fn assemble_async_iterator_candidates(
+        &mut self,
+        obligation: &PolyTraitObligation<'tcx>,
+        candidates: &mut SelectionCandidateSet<'tcx>,
+    ) {
+        let self_ty = obligation.self_ty().skip_binder();
+        if let ty::Coroutine(did, args, _) = *self_ty.kind() {
+            // gen constructs get lowered to a special kind of coroutine that
+            // should directly `impl AsyncIterator`.
+            if self.tcx().coroutine_is_async_gen(did) {
+                debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);
+
+                // Can only confirm this candidate if we have constrained
+                // the `Yield` type to at least `Poll<Option<?0>>`..
+                let ty::Adt(_poll_def, args) = *args.as_coroutine().yield_ty().kind() else {
+                    candidates.ambiguous = true;
+                    return;
+                };
+                let ty::Adt(_option_def, _) = *args.type_at(0).kind() else {
+                    candidates.ambiguous = true;
+                    return;
+                };
+
+                candidates.vec.push(AsyncIteratorCandidate);
+            }
+        }
+    }
+
     /// Checks for the artificial impl that the compiler will create for an obligation like `X :
     /// FnMut<..>` where `X` is a closure type.
     ///

compiler/rustc_trait_selection/src/traits/select/confirmation.rs

@@ -98,6 +98,11 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 ImplSource::Builtin(BuiltinImplSource::Misc, vtable_iterator)
             }
 
+            AsyncIteratorCandidate => {
+                let vtable_iterator = self.confirm_async_iterator_candidate(obligation)?;
+                ImplSource::Builtin(BuiltinImplSource::Misc, vtable_iterator)
+            }
+
             FnPointerCandidate { is_const } => {
                 let data = self.confirm_fn_pointer_candidate(obligation, is_const)?;
                 ImplSource::Builtin(BuiltinImplSource::Misc, data)
@@ -813,6 +818,35 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         Ok(nested)
     }
 
+    fn confirm_async_iterator_candidate(
+        &mut self,
+        obligation: &PolyTraitObligation<'tcx>,
+    ) -> Result<Vec<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
+        // Okay to skip binder because the args on coroutine types never
+        // touch bound regions, they just capture the in-scope
+        // type/region parameters.
+        let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder());
+        let ty::Coroutine(coroutine_def_id, args, _) = *self_ty.kind() else {
+            bug!("closure candidate for non-closure {:?}", obligation);
+        };
+
+        debug!(?obligation, ?coroutine_def_id, ?args, "confirm_async_iterator_candidate");
+
+        let gen_sig = args.as_coroutine().sig();
+
+        let (trait_ref, _) = super::util::async_iterator_trait_ref_and_outputs(
+            self.tcx(),
+            obligation.predicate.def_id(),
+            obligation.predicate.no_bound_vars().expect("iterator has no bound vars").self_ty(),
+            gen_sig,
+        );
+
+        let nested = self.confirm_poly_trait_refs(obligation, ty::Binder::dummy(trait_ref))?;
+        debug!(?trait_ref, ?nested, "iterator candidate obligations");
+
+        Ok(nested)
+    }
+
     #[instrument(skip(self), level = "debug")]
     fn confirm_closure_candidate(
         &mut self,

compiler/rustc_trait_selection/src/traits/select/mod.rs

@@ -1875,6 +1875,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate
@@ -1904,6 +1905,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate
@@ -1939,6 +1941,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate
@@ -1954,6 +1957,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate
@@ -2061,6 +2065,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate
@@ -2072,6 +2077,7 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
                 | CoroutineCandidate
                 | FutureCandidate
                 | IteratorCandidate
+                | AsyncIteratorCandidate
                 | FnPointerCandidate { .. }
                 | BuiltinObjectCandidate
                 | BuiltinUnsizeCandidate