wip: de-asyncfn-ify iter macro

Author: eholk

compiler/rustc_const_eval/src/check_consts/check.rs

@@ -591,7 +591,7 @@ impl<'tcx> Visitor<'tcx> for Checker<'_, 'tcx> {
                 if let AggregateKind::Coroutine(def_id, ..) = kind.as_ref()
                     && let Some(
                         coroutine_kind @ hir::CoroutineKind::Desugared(
-                            hir::CoroutineDesugaring::Async,
+                            hir::CoroutineDesugaring::Async | hir::CoroutineDesugaring::Gen,
                             _,
                         ),
                     ) = self.tcx.coroutine_kind(def_id)

compiler/rustc_hir_typeck/src/closure.rs

@@ -465,7 +465,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
             if let Some(trait_def_id) = trait_def_id {
                 let found_kind = match closure_kind {
-                    hir::ClosureKind::Closure => self.tcx.fn_trait_kind_from_def_id(trait_def_id),
+                    hir::ClosureKind::Closure
+                    // FIXME(iter_macro): Someday we'll probably want iterator closures instead of
+                    // just using Fn* for iterators.
+                    | hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Gen) => {
+                        self.tcx.fn_trait_kind_from_def_id(trait_def_id)
+                    }
                     hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Async) => self
                         .tcx
                         .async_fn_trait_kind_from_def_id(trait_def_id)

compiler/rustc_mir_transform/src/check_inline.rs

@@ -8,6 +8,7 @@ use rustc_middle::mir::{Body, TerminatorKind};
 use rustc_middle::ty;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::sym;
+use tracing::debug;
 
 use crate::pass_manager::MirLint;
 
@@ -40,7 +41,9 @@ pub(super) fn is_inline_valid_on_fn<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: DefId,
 ) -> Result<(), &'static str> {
+    debug!("is_inline_valid_on_fn({def_id:?})");
     let codegen_attrs = tcx.codegen_fn_attrs(def_id);
+    debug!("    codegen_attrs: {codegen_attrs:?}");
     if tcx.has_attr(def_id, sym::rustc_no_mir_inline) {
         return Err("#[rustc_no_mir_inline]");
     }

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

@@ -11,7 +11,7 @@ use std::ops::ControlFlow;
 use hir::LangItem;
 use hir::def_id::DefId;
 use rustc_data_structures::fx::{FxHashSet, FxIndexSet};
-use rustc_hir as hir;
+use rustc_hir::{self as hir, CoroutineDesugaring, CoroutineKind};
 use rustc_infer::traits::{Obligation, PolyTraitObligation, SelectionError};
 use rustc_middle::ty::fast_reject::DeepRejectCtxt;
 use rustc_middle::ty::{self, Ty, TypeVisitableExt, TypingMode};
@@ -128,11 +128,15 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     self.assemble_async_iterator_candidates(obligation, &mut candidates);
                 } else if tcx.is_lang_item(def_id, LangItem::AsyncFnKindHelper) {
                     self.assemble_async_fn_kind_helper_candidates(obligation, &mut candidates);
+                } else if tcx.is_lang_item(def_id, LangItem::AsyncFn)
+                    || tcx.is_lang_item(def_id, LangItem::AsyncFnOnce)
+                    || tcx.is_lang_item(def_id, LangItem::AsyncFnMut)
+                {
+                    self.assemble_async_closure_candidates(obligation, &mut candidates);
                 }
 
                 // FIXME: Put these into `else if` blocks above, since they're built-in.
                 self.assemble_closure_candidates(obligation, &mut candidates);
-                self.assemble_async_closure_candidates(obligation, &mut candidates);
                 self.assemble_fn_pointer_candidates(obligation, &mut candidates);
 
                 self.assemble_candidates_from_impls(obligation, &mut candidates);
@@ -428,6 +432,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         }
     }
 
+    #[instrument(level = "debug", skip(self, candidates))]
     fn assemble_async_closure_candidates(
         &mut self,
         obligation: &PolyTraitObligation<'tcx>,
@@ -439,15 +444,33 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             return;
         };
 
+        debug!("self_ty = {:?}", obligation.self_ty().skip_binder().kind());
         match *obligation.self_ty().skip_binder().kind() {
-            ty::CoroutineClosure(_, args) => {
+            ty::CoroutineClosure(def_id, args) => {
                 if let Some(closure_kind) =
                     args.as_coroutine_closure().kind_ty().to_opt_closure_kind()
                     && !closure_kind.extends(goal_kind)
                 {
                     return;
                 }
-                candidates.vec.push(AsyncClosureCandidate);
+
+                // Make sure this is actually an async closure.
+                let Some(coroutine_kind) =
+                    self.tcx().coroutine_kind(self.tcx().coroutine_for_closure(def_id))
+                else {
+                    bug!("coroutine with no kind");
+                };
+
+                debug!(?coroutine_kind);
+                match coroutine_kind {
+                    CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
+                        candidates.vec.push(AsyncClosureCandidate);
+                    }
+                    // CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
+                    //     candidates.vec.push(IteratorClosureCandidate)
+                    // }
+                    _ => (),
+                }
             }
             // Closures and fn pointers implement `AsyncFn*` if their return types
             // implement `Future`, which is checked later.

tests/ui/iterators/generator_capture_fnonce.rs

@@ -0,0 +1,32 @@
+//@ run-pass
+
+#![feature(iter_macro, yield_expr)]
+
+use std::iter::iter;
+
+fn main() {
+    let i = {
+        let s = String::new();
+        iter! { move || {
+            yield s.len();
+            for x in 5..10 {
+                yield x * 2;
+            }
+        }}
+    };
+    test_iterator(i);
+}
+
+/// Exercise the iterator in a separate function to ensure it's not capturing anything it shoudln't.
+fn test_iterator<I: Iterator<Item = usize>>(i: impl FnOnce() -> I) {
+    let mut i = i();
+    assert_eq!(i.next(), Some(0));
+    assert_eq!(i.next(), Some(10));
+    assert_eq!(i.next(), Some(12));
+    assert_eq!(i.next(), Some(14));
+    assert_eq!(i.next(), Some(16));
+    assert_eq!(i.next(), Some(18));
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+    assert_eq!(i.next(), None);
+}