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: [PATCH] Add additional tests and update existing tests

---
 compiler/rustc_parse/src/parser/expr.rs       |  3 +-
 .../ui/async-await/feature-async-for-loop.rs  |  9 ++++++
 .../async-await/feature-async-for-loop.stderr | 12 ++++++-
 tests/ui/async-await/for-await-2015.rs        | 10 ++++++
 .../ui/async-await/for-await-consumes-iter.rs | 20 ++++++++++++
 .../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

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) {}
+}