Add additional tests and update existing tests

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);

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

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`.

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

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

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`.

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