Add async_iter_next to assist with for await desugaring

Author: eholk

compiler/rustc_hir/src/lang_items.rs

@@ -307,6 +307,8 @@ language_item_table! {
     Context,                 sym::Context,             context,                    Target::Struct,         GenericRequirement::None;
     FuturePoll,              sym::poll,                future_poll_fn,             Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
 
+    AsyncIterNext,           sym::async_iter_next,     async_iter_next,            Target::Fn,             GenericRequirement::Exact(2);
+
     Option,                  sym::Option,              option_type,                Target::Enum,           GenericRequirement::None;
     OptionSome,              sym::Some,                option_some_variant,        Target::Variant,        GenericRequirement::None;
     OptionNone,              sym::None,                option_none_variant,        Target::Variant,        GenericRequirement::None;

compiler/rustc_span/src/symbol.rs

@@ -426,6 +426,7 @@ symbols! {
         async_closure,
         async_fn_in_trait,
         async_fn_track_caller,
+        async_iter_next,
         async_iterator,
         atomic,
         atomic_mod,

library/core/src/async_iter/async_iter.rs

@@ -156,3 +156,21 @@ impl<I: AsyncIterator> IntoAsyncIterator for I {
         self
     }
 }
+
+/// Create a future that returns the next item from an `AsyncIterator`.
+///
+/// This primarily exists for the compiler to use in desugaring `for await` loops.
+#[unstable(feature = "async_iterator", issue = "79024")]
+#[cfg_attr(not(bootstrap), lang = "async_iter_next")]
+pub fn async_iter_next<'a, I: AsyncIterator>(
+    iter: Pin<&'a mut I>,
+) -> impl core::future::Future<Output = Option<I::Item>> + 'a {
+    struct Next<'a, I: AsyncIterator>(Pin<&'a mut I>);
+    impl<'a, I: AsyncIterator> core::future::Future for Next<'a, I> {
+        type Output = Option<I::Item>;
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            I::poll_next(self.0.as_mut(), cx)
+        }
+    }
+    Next(iter)
+}

library/core/src/async_iter/mod.rs

@@ -124,5 +124,5 @@
 mod async_iter;
 mod from_iter;
 
-pub use async_iter::{AsyncIterator, IntoAsyncIterator};
+pub use async_iter::{async_iter_next, AsyncIterator, IntoAsyncIterator};
 pub use from_iter::{from_iter, FromIter};

library/core/tests/async_iter/mod.rs

@@ -1,4 +1,5 @@
 use core::async_iter::{self, AsyncIterator, IntoAsyncIterator};
+use core::future::Future;
 use core::pin::pin;
 use core::task::Poll;
 
@@ -25,3 +26,24 @@ fn into_async_iter() {
     assert_eq!(async_iter.as_mut().poll_next(&mut cx), Poll::Ready(Some(2)));
     assert_eq!(async_iter.as_mut().poll_next(&mut cx), Poll::Ready(None));
 }
+
+#[test]
+fn async_iter_next() {
+    let async_iter = async_iter::from_iter(0..3);
+    let async_iter = pin!(async_iter.into_async_iter());
+
+    let waker = noop_waker::NOOP_WAKER;
+    let mut cx = &mut core::task::Context::from_waker(&waker);
+
+    let mut fut = pin!(core::async_iter::async_iter_next(async_iter));
+
+    loop {
+        match fut.as_mut().poll(&mut cx) {
+            Poll::Ready(item) => {
+                assert_eq!(item, Some(0));
+                return;
+            }
+            Poll::Pending => continue,
+        }
+    }
+}