From 27d6539a46123dcdb6fae6e043b8c1c12b3e0d6f Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Fri, 8 Dec 2023 14:51:50 -0800
Subject: Plumb awaitness of for loops

---
 compiler/rustc_parse/src/parser/expr.rs | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

(limited to 'compiler/rustc_parse/src/parser')

diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index cd3e8b92f2f..f424040e4ea 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -10,7 +10,7 @@ use super::{
 use crate::errors;
 use crate::maybe_recover_from_interpolated_ty_qpath;
 use ast::mut_visit::{noop_visit_expr, MutVisitor};
-use ast::{CoroutineKind, GenBlockKind, Pat, Path, PathSegment};
+use ast::{CoroutineKind, ForLoopKind, GenBlockKind, Pat, Path, PathSegment};
 use core::mem;
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Delimiter, Token, TokenKind};
@@ -1801,7 +1801,7 @@ impl<'a> Parser<'a> {
                     && matches!(
                         expr.kind,
                         ExprKind::While(_, _, None)
-                            | ExprKind::ForLoop(_, _, _, None)
+                            | ExprKind::ForLoop { label: None, .. }
                             | ExprKind::Loop(_, None, _)
                             | ExprKind::Block(_, None)
                     )
@@ -2682,8 +2682,16 @@ impl<'a> Parser<'a> {
         Ok((pat, expr))
     }
 
-    /// Parses `for <src_pat> in <src_expr> <src_loop_block>` (`for` token already eaten).
+    /// Parses `for await? <src_pat> in <src_expr> <src_loop_block>` (`for` token already eaten).
     fn parse_expr_for(&mut self, opt_label: Option<Label>, lo: Span) -> PResult<'a, P<Expr>> {
+        let is_await = self.eat_keyword(kw::Await);
+
+        if is_await {
+            self.sess.gated_spans.gate(sym::async_for_loop, self.prev_token.span);
+        }
+
+        let kind = if is_await { ForLoopKind::ForAwait } else { ForLoopKind::For };
+
         let (pat, expr) = self.parse_for_head()?;
         // Recover from missing expression in `for` loop
         if matches!(expr.kind, ExprKind::Block(..))
@@ -2696,13 +2704,13 @@ impl<'a> Parser<'a> {
             let block = self.mk_block(thin_vec![], BlockCheckMode::Default, self.prev_token.span);
             return Ok(self.mk_expr(
                 lo.to(self.prev_token.span),
-                ExprKind::ForLoop(pat, err_expr, block, opt_label),
+                ExprKind::ForLoop { pat, iter: err_expr, body: block, label: opt_label, kind },
             ));
         }
 
         let (attrs, loop_block) = self.parse_inner_attrs_and_block()?;
 
-        let kind = ExprKind::ForLoop(pat, expr, loop_block, opt_label);
+        let kind = ExprKind::ForLoop { pat, iter: expr, body: loop_block, label: opt_label, kind };
 
         self.recover_loop_else("for", lo)?;
 
@@ -3798,7 +3806,7 @@ impl MutVisitor for CondChecker<'_> {
             | ExprKind::Lit(_)
             | ExprKind::If(_, _, _)
             | ExprKind::While(_, _, _)
-            | ExprKind::ForLoop(_, _, _, _)
+            | ExprKind::ForLoop { .. }
             | ExprKind::Loop(_, _, _)
             | ExprKind::Match(_, _)
             | ExprKind::Closure(_)
-- 
cgit 1.4.1-3-g733a5


From 397f4a15bbcad7d22f8c09edb72a0d49e3da80c9 Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Tue, 12 Dec 2023 17:33:51 -0800
Subject: Add additional tests and update existing tests

---
 compiler/rustc_parse/src/parser/expr.rs            |  3 +-
 tests/ui/async-await/feature-async-for-loop.rs     |  9 ++++++
 tests/ui/async-await/feature-async-for-loop.stderr | 12 +++++++-
 tests/ui/async-await/for-await-2015.rs             | 10 +++++++
 tests/ui/async-await/for-await-consumes-iter.rs    | 20 ++++++++++++++
 .../ui/async-await/for-await-consumes-iter.stderr  | 27 ++++++++++++++++++
 tests/ui/async-await/for-await-passthrough.rs      | 32 ++++++++++++++++++++++
 7 files changed, 111 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/async-await/for-await-2015.rs
 create mode 100644 tests/ui/async-await/for-await-consumes-iter.rs
 create mode 100644 tests/ui/async-await/for-await-consumes-iter.stderr
 create mode 100644 tests/ui/async-await/for-await-passthrough.rs

(limited to 'compiler/rustc_parse/src/parser')

diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index f424040e4ea..85f785b3401 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -2684,7 +2684,8 @@ impl<'a> Parser<'a> {
 
     /// Parses `for await? <src_pat> in <src_expr> <src_loop_block>` (`for` token already eaten).
     fn parse_expr_for(&mut self, opt_label: Option<Label>, lo: Span) -> PResult<'a, P<Expr>> {
-        let is_await = self.eat_keyword(kw::Await);
+        let is_await =
+            self.token.uninterpolated_span().at_least_rust_2018() && self.eat_keyword(kw::Await);
 
         if is_await {
             self.sess.gated_spans.gate(sym::async_for_loop, self.prev_token.span);
diff --git a/tests/ui/async-await/feature-async-for-loop.rs b/tests/ui/async-await/feature-async-for-loop.rs
index 6d32b3ca700..42247dd14b0 100644
--- a/tests/ui/async-await/feature-async-for-loop.rs
+++ b/tests/ui/async-await/feature-async-for-loop.rs
@@ -11,4 +11,13 @@ fn f() {
     };
 }
 
+#[cfg(FALSE)]
+fn g() {
+    let _ = async {
+        for await _i in core::async_iter::from_iter(0..3) {
+            //~^ ERROR `for await` loops are experimental
+        }
+    };
+}
+
 fn main() {}
diff --git a/tests/ui/async-await/feature-async-for-loop.stderr b/tests/ui/async-await/feature-async-for-loop.stderr
index 7e6d83fe5b6..38f75821772 100644
--- a/tests/ui/async-await/feature-async-for-loop.stderr
+++ b/tests/ui/async-await/feature-async-for-loop.stderr
@@ -4,8 +4,18 @@ error[E0658]: `for await` loops are experimental
 LL |         for await _i in core::async_iter::from_iter(0..3) {
    |             ^^^^^
    |
+   = note: see issue #118898 <https://github.com/rust-lang/rust/issues/118898> for more information
    = help: add `#![feature(async_for_loop)]` to the crate attributes to enable
 
-error: aborting due to 1 previous error
+error[E0658]: `for await` loops are experimental
+  --> $DIR/feature-async-for-loop.rs:17:13
+   |
+LL |         for await _i in core::async_iter::from_iter(0..3) {
+   |             ^^^^^
+   |
+   = note: see issue #118898 <https://github.com/rust-lang/rust/issues/118898> for more information
+   = help: add `#![feature(async_for_loop)]` to the crate attributes to enable
+
+error: aborting due to 2 previous errors
 
 For more information about this error, try `rustc --explain E0658`.
diff --git a/tests/ui/async-await/for-await-2015.rs b/tests/ui/async-await/for-await-2015.rs
new file mode 100644
index 00000000000..c1b7c016d1f
--- /dev/null
+++ b/tests/ui/async-await/for-await-2015.rs
@@ -0,0 +1,10 @@
+// check-pass
+
+#![feature(async_for_loop)]
+
+// Make sure we don't break `for await` loops in the 2015 edition, where `await` was allowed as an
+// identifier.
+
+fn main() {
+    for await in 0..3 {}
+}
diff --git a/tests/ui/async-await/for-await-consumes-iter.rs b/tests/ui/async-await/for-await-consumes-iter.rs
new file mode 100644
index 00000000000..65bb9e88448
--- /dev/null
+++ b/tests/ui/async-await/for-await-consumes-iter.rs
@@ -0,0 +1,20 @@
+// edition: 2021
+#![feature(async_iterator, async_iter_from_iter, const_waker, async_for_loop, noop_waker)]
+
+use std::future::Future;
+
+// a test to make sure `for await` consumes the iterator
+
+async fn real_main() {
+    let iter = core::async_iter::from_iter(0..3);
+    let mut count = 0;
+    for await i in iter {
+    }
+    // make sure iter has been moved and we can't iterate over it again.
+    for await i in iter {
+        //~^ ERROR: use of moved value: `iter`
+    }
+}
+
+fn main() {
+}
diff --git a/tests/ui/async-await/for-await-consumes-iter.stderr b/tests/ui/async-await/for-await-consumes-iter.stderr
new file mode 100644
index 00000000000..48b2e8a51d8
--- /dev/null
+++ b/tests/ui/async-await/for-await-consumes-iter.stderr
@@ -0,0 +1,27 @@
+error[E0382]: use of moved value: `iter`
+  --> $DIR/for-await-consumes-iter.rs:14:20
+   |
+LL |     let iter = core::async_iter::from_iter(0..3);
+   |         ---- move occurs because `iter` has type `FromIter<std::ops::Range<i32>>`, which does not implement the `Copy` trait
+LL |     let mut count = 0;
+LL |     for await i in iter {
+   |     -------------------
+   |     |              |
+   |     |              value moved here
+   |     inside of this loop
+...
+LL |     for await i in iter {
+   |                    ^^^^ value used here after move
+   |
+help: consider cloning the value if the performance cost is acceptable
+   |
+LL |     for await i in iter.clone() {
+   |                        ++++++++
+help: borrow this binding in the pattern to avoid moving the value
+   |
+LL |     for await i in ref iter {
+   |                    +++
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0382`.
diff --git a/tests/ui/async-await/for-await-passthrough.rs b/tests/ui/async-await/for-await-passthrough.rs
new file mode 100644
index 00000000000..7fa133aaedc
--- /dev/null
+++ b/tests/ui/async-await/for-await-passthrough.rs
@@ -0,0 +1,32 @@
+// run-pass
+// edition: 2024
+// compile-flags: -Zunstable-options
+#![feature(async_iterator, async_iter_from_iter, const_waker, async_for_loop, noop_waker,
+           gen_blocks)]
+
+use std::future::Future;
+
+async gen fn async_iter() -> i32 {
+    let iter = core::async_iter::from_iter(0..3);
+    for await i in iter {
+        yield i + 1;
+    }
+}
+
+// make sure a simple for await loop works
+async fn real_main() {
+    let mut count = 1;
+    for await i in async_iter() {
+        assert_eq!(i, count);
+        count += 1;
+    }
+    assert_eq!(count, 4);
+}
+
+fn main() {
+    let future = real_main();
+    let waker = std::task::Waker::noop();
+    let mut cx = &mut core::task::Context::from_waker(&waker);
+    let mut future = core::pin::pin!(future);
+    while let core::task::Poll::Pending = future.as_mut().poll(&mut cx) {}
+}
-- 
cgit 1.4.1-3-g733a5