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: [PATCH 1/3] Plumb awaitness of for loops

---
 compiler/rustc_ast/src/ast.rs                 | 13 +++++++---
 compiler/rustc_ast/src/mut_visit.rs           |  2 +-
 compiler/rustc_ast/src/util/classify.rs       | 14 ++++++++---
 compiler/rustc_ast/src/visit.rs               | 10 ++++----
 compiler/rustc_ast_lowering/src/expr.rs       |  9 +++----
 compiler/rustc_ast_passes/src/feature_gate.rs |  1 +
 .../rustc_ast_pretty/src/pprust/state/expr.rs | 10 +++++---
 .../src/assert/context.rs                     |  2 +-
 compiler/rustc_feature/src/unstable.rs        |  2 ++
 compiler/rustc_lint/src/unused.rs             |  6 ++---
 compiler/rustc_parse/src/parser/expr.rs       | 20 +++++++++++-----
 compiler/rustc_resolve/src/late.rs            |  6 ++---
 .../rustc_resolve/src/late/diagnostics.rs     |  2 +-
 compiler/rustc_span/src/symbol.rs             |  1 +
 .../clippy_lints/src/needless_continue.rs     |  6 ++++-
 .../clippy/clippy_lints/src/redundant_else.rs |  2 +-
 .../src/suspicious_operation_groupings.rs     |  2 +-
 .../clippy/clippy_utils/src/ast_utils.rs      | 19 ++++++++++++---
 src/tools/clippy/clippy_utils/src/sugg.rs     |  2 +-
 src/tools/rustfmt/src/closures.rs             |  4 ++--
 src/tools/rustfmt/src/expr.rs                 | 24 +++++++++++++------
 src/tools/rustfmt/src/matches.rs              |  2 +-
 src/tools/rustfmt/src/overflow.rs             |  2 +-
 src/tools/rustfmt/src/utils.rs                |  4 ++--
 .../ui/async-await/feature-async-for-loop.rs  | 14 +++++++++++
 .../async-await/feature-async-for-loop.stderr | 11 +++++++++
 26 files changed, 137 insertions(+), 53 deletions(-)
 create mode 100644 tests/ui/async-await/feature-async-for-loop.rs
 create mode 100644 tests/ui/async-await/feature-async-for-loop.stderr

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 9d543563c0f96..8a0fbdd1357ae 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -1274,7 +1274,7 @@ impl Expr {
             ExprKind::Let(..) => ExprPrecedence::Let,
             ExprKind::If(..) => ExprPrecedence::If,
             ExprKind::While(..) => ExprPrecedence::While,
-            ExprKind::ForLoop(..) => ExprPrecedence::ForLoop,
+            ExprKind::ForLoop { .. } => ExprPrecedence::ForLoop,
             ExprKind::Loop(..) => ExprPrecedence::Loop,
             ExprKind::Match(..) => ExprPrecedence::Match,
             ExprKind::Closure(..) => ExprPrecedence::Closure,
@@ -1436,10 +1436,10 @@ pub enum ExprKind {
     While(P<Expr>, P<Block>, Option<Label>),
     /// A `for` loop, with an optional label.
     ///
-    /// `'label: for pat in expr { block }`
+    /// `'label: for await? pat in iter { block }`
     ///
     /// This is desugared to a combination of `loop` and `match` expressions.
-    ForLoop(P<Pat>, P<Expr>, P<Block>, Option<Label>),
+    ForLoop { pat: P<Pat>, iter: P<Expr>, body: P<Block>, label: Option<Label>, kind: ForLoopKind },
     /// Conditionless loop (can be exited with `break`, `continue`, or `return`).
     ///
     /// `'label: loop { block }`
@@ -1542,6 +1542,13 @@ pub enum ExprKind {
     Err,
 }
 
+/// Used to differentiate between `for` loops and `for await` loops.
+#[derive(Clone, Copy, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub enum ForLoopKind {
+    For,
+    ForAwait,
+}
+
 /// Used to differentiate between `async {}` blocks and `gen {}` blocks.
 #[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
 pub enum GenBlockKind {
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 10b2025f93783..d0afd90e5e75b 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1389,7 +1389,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             vis.visit_block(body);
             visit_opt(label, |label| vis.visit_label(label));
         }
-        ExprKind::ForLoop(pat, iter, body, label) => {
+        ExprKind::ForLoop { pat, iter, body, label, kind: _ } => {
             vis.visit_pat(pat);
             vis.visit_expr(iter);
             vis.visit_block(body);
diff --git a/compiler/rustc_ast/src/util/classify.rs b/compiler/rustc_ast/src/util/classify.rs
index 4dece07978376..65036bcdc36fd 100644
--- a/compiler/rustc_ast/src/util/classify.rs
+++ b/compiler/rustc_ast/src/util/classify.rs
@@ -19,7 +19,7 @@ pub fn expr_requires_semi_to_be_stmt(e: &ast::Expr) -> bool {
             | ast::ExprKind::Block(..)
             | ast::ExprKind::While(..)
             | ast::ExprKind::Loop(..)
-            | ast::ExprKind::ForLoop(..)
+            | ast::ExprKind::ForLoop { .. }
             | ast::ExprKind::TryBlock(..)
             | ast::ExprKind::ConstBlock(..)
     )
@@ -48,8 +48,16 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<&ast::Expr> {
             Closure(closure) => {
                 expr = &closure.body;
             }
-            Gen(..) | Block(..) | ForLoop(..) | If(..) | Loop(..) | Match(..) | Struct(..)
-            | TryBlock(..) | While(..) | ConstBlock(_) => break Some(expr),
+            Gen(..)
+            | Block(..)
+            | ForLoop { .. }
+            | If(..)
+            | Loop(..)
+            | Match(..)
+            | Struct(..)
+            | TryBlock(..)
+            | While(..)
+            | ConstBlock(_) => break Some(expr),
 
             // FIXME: These can end in `}`, but changing these would break stable code.
             InlineAsm(_) | OffsetOf(_, _) | MacCall(_) | IncludedBytes(_) | FormatArgs(_) => {
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index 27f1b84f37251..45261ca48fc25 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -844,11 +844,11 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             visitor.visit_expr(subexpression);
             visitor.visit_block(block);
         }
-        ExprKind::ForLoop(pattern, subexpression, block, opt_label) => {
-            walk_list!(visitor, visit_label, opt_label);
-            visitor.visit_pat(pattern);
-            visitor.visit_expr(subexpression);
-            visitor.visit_block(block);
+        ExprKind::ForLoop { pat, iter, body, label, kind: _ } => {
+            walk_list!(visitor, visit_label, label);
+            visitor.visit_pat(pat);
+            visitor.visit_expr(iter);
+            visitor.visit_block(body);
         }
         ExprKind::Loop(block, opt_label, _) => {
             walk_list!(visitor, visit_label, opt_label);
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 11b5131b8d788..c2d91b0945304 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -56,12 +56,12 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     return ex;
                 }
                 // Desugar `ExprForLoop`
-                // from: `[opt_ident]: for <pat> in <head> <body>`
+                // from: `[opt_ident]: for await? <pat> in <iter> <body>`
                 //
                 // This also needs special handling because the HirId of the returned `hir::Expr` will not
                 // correspond to the `e.id`, so `lower_expr_for` handles attribute lowering itself.
-                ExprKind::ForLoop(pat, head, body, opt_label) => {
-                    return self.lower_expr_for(e, pat, head, body, *opt_label);
+                ExprKind::ForLoop { pat, iter, body, label, kind } => {
+                    return self.lower_expr_for(e, pat, iter, body, *label, *kind);
                 }
                 _ => (),
             }
@@ -337,7 +337,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ),
                 ExprKind::Try(sub_expr) => self.lower_expr_try(e.span, sub_expr),
 
-                ExprKind::Paren(_) | ExprKind::ForLoop(..) => {
+                ExprKind::Paren(_) | ExprKind::ForLoop{..} => {
                     unreachable!("already handled")
                 }
 
@@ -1673,6 +1673,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         head: &Expr,
         body: &Block,
         opt_label: Option<Label>,
+        _loop_kind: ForLoopKind,
     ) -> hir::Expr<'hir> {
         let head = self.lower_expr_mut(head);
         let pat = self.lower_pat(pat);
diff --git a/compiler/rustc_ast_passes/src/feature_gate.rs b/compiler/rustc_ast_passes/src/feature_gate.rs
index 142cdd15e64ac..411c10df0c21e 100644
--- a/compiler/rustc_ast_passes/src/feature_gate.rs
+++ b/compiler/rustc_ast_passes/src/feature_gate.rs
@@ -526,6 +526,7 @@ pub fn check_crate(krate: &ast::Crate, sess: &Session, features: &Features) {
         "async closures are unstable",
         "to use an async block, remove the `||`: `async {`"
     );
+    gate_all!(async_for_loop, "`for await` loops are experimental");
     gate_all!(
         closure_lifetime_binder,
         "`for<...>` binders for closures are experimental",
diff --git a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
index f5ffcddb83d74..f868beec81221 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
@@ -1,5 +1,6 @@
 use crate::pp::Breaks::Inconsistent;
 use crate::pprust::state::{AnnNode, PrintState, State, INDENT_UNIT};
+use ast::ForLoopKind;
 use itertools::{Itertools, Position};
 use rustc_ast::ptr::P;
 use rustc_ast::token;
@@ -418,20 +419,23 @@ impl<'a> State<'a> {
                 self.space();
                 self.print_block_with_attrs(blk, attrs);
             }
-            ast::ExprKind::ForLoop(pat, iter, blk, opt_label) => {
-                if let Some(label) = opt_label {
+            ast::ExprKind::ForLoop { pat, iter, body, label, kind } => {
+                if let Some(label) = label {
                     self.print_ident(label.ident);
                     self.word_space(":");
                 }
                 self.cbox(0);
                 self.ibox(0);
                 self.word_nbsp("for");
+                if kind == &ForLoopKind::ForAwait {
+                    self.word_nbsp("await");
+                }
                 self.print_pat(pat);
                 self.space();
                 self.word_space("in");
                 self.print_expr_as_cond(iter);
                 self.space();
-                self.print_block_with_attrs(blk, attrs);
+                self.print_block_with_attrs(body, attrs);
             }
             ast::ExprKind::Loop(blk, opt_label, _) => {
                 if let Some(label) = opt_label {
diff --git a/compiler/rustc_builtin_macros/src/assert/context.rs b/compiler/rustc_builtin_macros/src/assert/context.rs
index b54e119189abb..4e07122a5c4c9 100644
--- a/compiler/rustc_builtin_macros/src/assert/context.rs
+++ b/compiler/rustc_builtin_macros/src/assert/context.rs
@@ -303,7 +303,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::Continue(_)
             | ExprKind::Err
             | ExprKind::Field(_, _)
-            | ExprKind::ForLoop(_, _, _, _)
+            | ExprKind::ForLoop {..}
             | ExprKind::FormatArgs(_)
             | ExprKind::IncludedBytes(..)
             | ExprKind::InlineAsm(_)
diff --git a/compiler/rustc_feature/src/unstable.rs b/compiler/rustc_feature/src/unstable.rs
index bbf5e03117572..fdd4bb755e415 100644
--- a/compiler/rustc_feature/src/unstable.rs
+++ b/compiler/rustc_feature/src/unstable.rs
@@ -357,6 +357,8 @@ declare_features! (
     (unstable, async_closure, "1.37.0", Some(62290)),
     /// Allows `#[track_caller]` on async functions.
     (unstable, async_fn_track_caller, "1.73.0", Some(110011)),
+    /// Allows `for await` loops.
+    (unstable, async_for_loop, "CURRENT_RUSTC_VERSION", None),
     /// Allows builtin # foo() syntax
     (unstable, builtin_syntax, "1.71.0", Some(110680)),
     /// Treat `extern "C"` function as nounwind.
diff --git a/compiler/rustc_lint/src/unused.rs b/compiler/rustc_lint/src/unused.rs
index 34cdee4ec5c1a..0386f2ec56c0a 100644
--- a/compiler/rustc_lint/src/unused.rs
+++ b/compiler/rustc_lint/src/unused.rs
@@ -852,8 +852,8 @@ trait UnusedDelimLint {
                 (cond, UnusedDelimsCtx::WhileCond, true, Some(left), Some(right), true)
             }
 
-            ForLoop(_, ref cond, ref block, ..) => {
-                (cond, UnusedDelimsCtx::ForIterExpr, true, None, Some(block.span.lo()), true)
+            ForLoop { ref iter, ref body, .. } => {
+                (iter, UnusedDelimsCtx::ForIterExpr, true, None, Some(body.span.lo()), true)
             }
 
             Match(ref head, _) if Self::LINT_EXPR_IN_PATTERN_MATCHING_CTX => {
@@ -1085,7 +1085,7 @@ impl EarlyLintPass for UnusedParens {
         }
 
         match e.kind {
-            ExprKind::Let(ref pat, _, _, _) | ExprKind::ForLoop(ref pat, ..) => {
+            ExprKind::Let(ref pat, _, _, _) | ExprKind::ForLoop { ref pat, .. } => {
                 self.check_unused_parens_pat(cx, pat, false, false, (true, true));
             }
             // We ignore parens in cases like `if (((let Some(0) = Some(1))))` because we already
diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index cd3e8b92f2f9e..f424040e4ea7e 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(_)
diff --git a/compiler/rustc_resolve/src/late.rs b/compiler/rustc_resolve/src/late.rs
index 037179350f044..7e275a207e779 100644
--- a/compiler/rustc_resolve/src/late.rs
+++ b/compiler/rustc_resolve/src/late.rs
@@ -4252,11 +4252,11 @@ impl<'a: 'ast, 'b, 'ast, 'tcx> LateResolutionVisitor<'a, 'b, 'ast, 'tcx> {
                 });
             }
 
-            ExprKind::ForLoop(ref pat, ref iter_expr, ref block, label) => {
-                self.visit_expr(iter_expr);
+            ExprKind::ForLoop { ref pat, ref iter, ref body, label, kind: _ } => {
+                self.visit_expr(iter);
                 self.with_rib(ValueNS, RibKind::Normal, |this| {
                     this.resolve_pattern_top(pat, PatternSource::For);
-                    this.resolve_labeled_block(label, expr.id, block);
+                    this.resolve_labeled_block(label, expr.id, body);
                 });
             }
 
diff --git a/compiler/rustc_resolve/src/late/diagnostics.rs b/compiler/rustc_resolve/src/late/diagnostics.rs
index d767ed74139ba..b4dc9d35b423e 100644
--- a/compiler/rustc_resolve/src/late/diagnostics.rs
+++ b/compiler/rustc_resolve/src/late/diagnostics.rs
@@ -1409,7 +1409,7 @@ impl<'a: 'ast, 'ast, 'tcx> LateResolutionVisitor<'a, '_, 'ast, 'tcx> {
                             | ExprKind::Unary(..)
                             | ExprKind::If(..)
                             | ExprKind::While(..)
-                            | ExprKind::ForLoop(..)
+                            | ExprKind::ForLoop { .. }
                             | ExprKind::Match(..),
                         ..
                     }),
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 0333b5f04c3bc..2e7037532b5c6 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -426,6 +426,7 @@ symbols! {
         async_closure,
         async_fn_in_trait,
         async_fn_track_caller,
+        async_for_loop,
         async_iterator,
         atomic,
         atomic_mod,
diff --git a/src/tools/clippy/clippy_lints/src/needless_continue.rs b/src/tools/clippy/clippy_lints/src/needless_continue.rs
index 4b9ab50e4fd77..ff72b5e69ef1b 100644
--- a/src/tools/clippy/clippy_lints/src/needless_continue.rs
+++ b/src/tools/clippy/clippy_lints/src/needless_continue.rs
@@ -220,7 +220,11 @@ where
     F: FnMut(&ast::Block, Option<&ast::Label>),
 {
     if let ast::ExprKind::While(_, loop_block, label)
-    | ast::ExprKind::ForLoop(_, _, loop_block, label)
+    | ast::ExprKind::ForLoop {
+        body: loop_block,
+        label,
+        ..
+    }
     | ast::ExprKind::Loop(loop_block, label, ..) = &expr.kind
     {
         func(loop_block, label.as_ref());
diff --git a/src/tools/clippy/clippy_lints/src/redundant_else.rs b/src/tools/clippy/clippy_lints/src/redundant_else.rs
index 001686c84f818..fb434fb7450a7 100644
--- a/src/tools/clippy/clippy_lints/src/redundant_else.rs
+++ b/src/tools/clippy/clippy_lints/src/redundant_else.rs
@@ -111,7 +111,7 @@ impl<'ast> Visitor<'ast> for BreakVisitor {
             ExprKind::If(_, ref then, Some(ref els)) => self.check_block(then) && self.check_expr(els),
             ExprKind::If(_, _, None)
             // ignore loops for simplicity
-            | ExprKind::While(..) | ExprKind::ForLoop(..) | ExprKind::Loop(..) => false,
+            | ExprKind::While(..) | ExprKind::ForLoop { .. } | ExprKind::Loop(..) => false,
             _ => {
                 walk_expr(self, expr);
                 return;
diff --git a/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs b/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
index 8b9d9bade915f..60e9d262e7e09 100644
--- a/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
+++ b/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
@@ -554,7 +554,7 @@ fn ident_difference_expr_with_base_location(
         | (Closure(_), Closure(_))
         | (Match(_, _), Match(_, _))
         | (Loop(_, _, _), Loop(_, _, _))
-        | (ForLoop(_, _, _, _), ForLoop(_, _, _, _))
+        | (ForLoop { .. }, ForLoop { .. })
         | (While(_, _, _), While(_, _, _))
         | (If(_, _, _), If(_, _, _))
         | (Let(_, _, _, _), Let(_, _, _, _))
diff --git a/src/tools/clippy/clippy_utils/src/ast_utils.rs b/src/tools/clippy/clippy_utils/src/ast_utils.rs
index e36f2fa87a72a..be3c418ef7076 100644
--- a/src/tools/clippy/clippy_utils/src/ast_utils.rs
+++ b/src/tools/clippy/clippy_utils/src/ast_utils.rs
@@ -169,9 +169,22 @@ pub fn eq_expr(l: &Expr, r: &Expr) -> bool {
         (Let(lp, le, _, _), Let(rp, re, _, _)) => eq_pat(lp, rp) && eq_expr(le, re),
         (If(lc, lt, le), If(rc, rt, re)) => eq_expr(lc, rc) && eq_block(lt, rt) && eq_expr_opt(le, re),
         (While(lc, lt, ll), While(rc, rt, rl)) => eq_label(ll, rl) && eq_expr(lc, rc) && eq_block(lt, rt),
-        (ForLoop(lp, li, lt, ll), ForLoop(rp, ri, rt, rl)) => {
-            eq_label(ll, rl) && eq_pat(lp, rp) && eq_expr(li, ri) && eq_block(lt, rt)
-        },
+        (
+            ForLoop {
+                pat: lp,
+                iter: li,
+                body: lt,
+                label: ll,
+                kind: lk,
+            },
+            ForLoop {
+                pat: rp,
+                iter: ri,
+                body: rt,
+                label: rl,
+                kind: rk,
+            },
+        ) => eq_label(ll, rl) && eq_pat(lp, rp) && eq_expr(li, ri) && eq_block(lt, rt) && lk == rk,
         (Loop(lt, ll, _), Loop(rt, rl, _)) => eq_label(ll, rl) && eq_block(lt, rt),
         (Block(lb, ll), Block(rb, rl)) => eq_label(ll, rl) && eq_block(lb, rb),
         (TryBlock(l), TryBlock(r)) => eq_block(l, r),
diff --git a/src/tools/clippy/clippy_utils/src/sugg.rs b/src/tools/clippy/clippy_utils/src/sugg.rs
index 9b2bc8df1f307..c86362c427ce2 100644
--- a/src/tools/clippy/clippy_utils/src/sugg.rs
+++ b/src/tools/clippy/clippy_utils/src/sugg.rs
@@ -197,7 +197,7 @@ impl<'a> Sugg<'a> {
             | ast::ExprKind::Continue(..)
             | ast::ExprKind::Yield(..)
             | ast::ExprKind::Field(..)
-            | ast::ExprKind::ForLoop(..)
+            | ast::ExprKind::ForLoop { .. }
             | ast::ExprKind::Index(..)
             | ast::ExprKind::InlineAsm(..)
             | ast::ExprKind::OffsetOf(..)
diff --git a/src/tools/rustfmt/src/closures.rs b/src/tools/rustfmt/src/closures.rs
index f698f494ae538..5bf29441b5421 100644
--- a/src/tools/rustfmt/src/closures.rs
+++ b/src/tools/rustfmt/src/closures.rs
@@ -448,7 +448,7 @@ fn is_block_closure_forced(context: &RewriteContext<'_>, expr: &ast::Expr) -> bo
 
 fn is_block_closure_forced_inner(expr: &ast::Expr, version: Version) -> bool {
     match expr.kind {
-        ast::ExprKind::If(..) | ast::ExprKind::While(..) | ast::ExprKind::ForLoop(..) => true,
+        ast::ExprKind::If(..) | ast::ExprKind::While(..) | ast::ExprKind::ForLoop { .. } => true,
         ast::ExprKind::Loop(..) if version == Version::Two => true,
         ast::ExprKind::AddrOf(_, _, ref expr)
         | ast::ExprKind::Try(ref expr)
@@ -473,7 +473,7 @@ fn expr_requires_semi_to_be_stmt(e: &ast::Expr) -> bool {
         | ast::ExprKind::Block(..)
         | ast::ExprKind::While(..)
         | ast::ExprKind::Loop(..)
-        | ast::ExprKind::ForLoop(..)
+        | ast::ExprKind::ForLoop { .. }
         | ast::ExprKind::TryBlock(..) => false,
         _ => true,
     }
diff --git a/src/tools/rustfmt/src/expr.rs b/src/tools/rustfmt/src/expr.rs
index a68bd6694ba62..7808f89133649 100644
--- a/src/tools/rustfmt/src/expr.rs
+++ b/src/tools/rustfmt/src/expr.rs
@@ -3,7 +3,7 @@ use std::cmp::min;
 
 use itertools::Itertools;
 use rustc_ast::token::{Delimiter, Lit, LitKind};
-use rustc_ast::{ast, ptr, token};
+use rustc_ast::{ast, ptr, token, ForLoopKind};
 use rustc_span::{BytePos, Span};
 
 use crate::chains::rewrite_chain;
@@ -134,7 +134,7 @@ pub(crate) fn format_expr(
         }
         ast::ExprKind::Let(ref pat, ref expr, _span, _) => rewrite_let(context, shape, pat, expr),
         ast::ExprKind::If(..)
-        | ast::ExprKind::ForLoop(..)
+        | ast::ExprKind::ForLoop { .. }
         | ast::ExprKind::Loop(..)
         | ast::ExprKind::While(..) => to_control_flow(expr, expr_type)
             .and_then(|control_flow| control_flow.rewrite(context, shape)),
@@ -682,9 +682,15 @@ fn to_control_flow(expr: &ast::Expr, expr_type: ExprType) -> Option<ControlFlow<
                 expr.span,
             ))
         }
-        ast::ExprKind::ForLoop(ref pat, ref cond, ref block, label) => {
-            Some(ControlFlow::new_for(pat, cond, block, label, expr.span))
-        }
+        ast::ExprKind::ForLoop {
+            ref pat,
+            ref iter,
+            ref body,
+            label,
+            kind,
+        } => Some(ControlFlow::new_for(
+            pat, iter, body, label, expr.span, kind,
+        )),
         ast::ExprKind::Loop(ref block, label, _) => {
             Some(ControlFlow::new_loop(block, label, expr.span))
         }
@@ -771,6 +777,7 @@ impl<'a> ControlFlow<'a> {
         block: &'a ast::Block,
         label: Option<ast::Label>,
         span: Span,
+        kind: ForLoopKind,
     ) -> ControlFlow<'a> {
         ControlFlow {
             cond: Some(cond),
@@ -778,7 +785,10 @@ impl<'a> ControlFlow<'a> {
             else_block: None,
             label,
             pat: Some(pat),
-            keyword: "for",
+            keyword: match kind {
+                ForLoopKind::For => "for",
+                ForLoopKind::ForAwait => "for await",
+            },
             matcher: "",
             connector: " in",
             allow_single_line: false,
@@ -1364,7 +1374,7 @@ pub(crate) fn can_be_overflowed_expr(
                 || context.config.overflow_delimited_expr()
         }
         ast::ExprKind::If(..)
-        | ast::ExprKind::ForLoop(..)
+        | ast::ExprKind::ForLoop { .. }
         | ast::ExprKind::Loop(..)
         | ast::ExprKind::While(..) => {
             context.config.combine_control_expr() && context.use_block_indent() && args_len == 1
diff --git a/src/tools/rustfmt/src/matches.rs b/src/tools/rustfmt/src/matches.rs
index ef509b568372d..5a00984d4c00e 100644
--- a/src/tools/rustfmt/src/matches.rs
+++ b/src/tools/rustfmt/src/matches.rs
@@ -591,7 +591,7 @@ fn can_flatten_block_around_this(body: &ast::Expr) -> bool {
         ast::ExprKind::If(..) => false,
         // We do not allow collapsing a block around expression with condition
         // to avoid it being cluttered with match arm.
-        ast::ExprKind::ForLoop(..) | ast::ExprKind::While(..) => false,
+        ast::ExprKind::ForLoop { .. } | ast::ExprKind::While(..) => false,
         ast::ExprKind::Loop(..)
         | ast::ExprKind::Match(..)
         | ast::ExprKind::Block(..)
diff --git a/src/tools/rustfmt/src/overflow.rs b/src/tools/rustfmt/src/overflow.rs
index d81bf24dbd1cd..f46583b1c8995 100644
--- a/src/tools/rustfmt/src/overflow.rs
+++ b/src/tools/rustfmt/src/overflow.rs
@@ -409,7 +409,7 @@ impl<'a> Context<'a> {
                     // When overflowing the expressions which consists of a control flow
                     // expression, avoid condition to use multi line.
                     ast::ExprKind::If(..)
-                    | ast::ExprKind::ForLoop(..)
+                    | ast::ExprKind::ForLoop { .. }
                     | ast::ExprKind::Loop(..)
                     | ast::ExprKind::While(..)
                     | ast::ExprKind::Match(..) => {
diff --git a/src/tools/rustfmt/src/utils.rs b/src/tools/rustfmt/src/utils.rs
index 7d7bbf1152905..642b6603b1e9a 100644
--- a/src/tools/rustfmt/src/utils.rs
+++ b/src/tools/rustfmt/src/utils.rs
@@ -295,7 +295,7 @@ pub(crate) fn semicolon_for_stmt(
 ) -> bool {
     match stmt.kind {
         ast::StmtKind::Semi(ref expr) => match expr.kind {
-            ast::ExprKind::While(..) | ast::ExprKind::Loop(..) | ast::ExprKind::ForLoop(..) => {
+            ast::ExprKind::While(..) | ast::ExprKind::Loop(..) | ast::ExprKind::ForLoop { .. } => {
                 false
             }
             ast::ExprKind::Break(..) | ast::ExprKind::Continue(..) | ast::ExprKind::Ret(..) => {
@@ -476,7 +476,7 @@ pub(crate) fn is_block_expr(context: &RewriteContext<'_>, expr: &ast::Expr, repr
         | ast::ExprKind::ConstBlock(..)
         | ast::ExprKind::Gen(..)
         | ast::ExprKind::Loop(..)
-        | ast::ExprKind::ForLoop(..)
+        | ast::ExprKind::ForLoop { .. }
         | ast::ExprKind::TryBlock(..)
         | ast::ExprKind::Match(..) => repr.contains('\n'),
         ast::ExprKind::Paren(ref expr)
diff --git a/tests/ui/async-await/feature-async-for-loop.rs b/tests/ui/async-await/feature-async-for-loop.rs
new file mode 100644
index 0000000000000..6d32b3ca7007b
--- /dev/null
+++ b/tests/ui/async-await/feature-async-for-loop.rs
@@ -0,0 +1,14 @@
+// edition:2021
+// gate-test-async_for_loop
+
+#![feature(async_iter_from_iter, async_iterator)]
+
+fn f() {
+    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
new file mode 100644
index 0000000000000..7e6d83fe5b6c8
--- /dev/null
+++ b/tests/ui/async-await/feature-async-for-loop.stderr
@@ -0,0 +1,11 @@
+error[E0658]: `for await` loops are experimental
+  --> $DIR/feature-async-for-loop.rs:8:13
+   |
+LL |         for await _i in core::async_iter::from_iter(0..3) {
+   |             ^^^^^
+   |
+   = help: add `#![feature(async_for_loop)]` to the crate attributes to enable
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0658`.

From 97df0d3657fafac6e6ca481c1d4123bf6e2633f1 Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Fri, 8 Dec 2023 17:00:11 -0800
Subject: [PATCH 2/3] Desugar for await loops

---
 compiler/rustc_ast_lowering/src/expr.rs       | 120 ++++++++++++++----
 compiler/rustc_ast_lowering/src/lib.rs        |   2 +
 .../src/assert/context.rs                     |   2 +-
 compiler/rustc_feature/src/unstable.rs        |   2 +-
 compiler/rustc_hir/src/lang_items.rs          |   2 +
 compiler/rustc_span/src/symbol.rs             |   2 +
 library/core/src/async_iter/async_iter.rs     |   1 +
 tests/ui/async-await/for-await.rs             |  24 ++++
 8 files changed, 125 insertions(+), 30 deletions(-)
 create mode 100644 tests/ui/async-await/for-await.rs

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index c2d91b0945304..660cc9b70f1c7 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -337,7 +337,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ),
                 ExprKind::Try(sub_expr) => self.lower_expr_try(e.span, sub_expr),
 
-                ExprKind::Paren(_) | ExprKind::ForLoop{..} => {
+                ExprKind::Paren(_) | ExprKind::ForLoop { .. } => {
                     unreachable!("already handled")
                 }
 
@@ -874,6 +874,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// }
     /// ```
     fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
+        let expr = self.arena.alloc(self.lower_expr_mut(expr));
+        self.make_lowered_await(await_kw_span, expr, FutureKind::Future)
+    }
+
+    /// Takes an expr that has already been lowered and generates a desugared await loop around it
+    fn make_lowered_await(
+        &mut self,
+        await_kw_span: Span,
+        expr: &'hir hir::Expr<'hir>,
+        await_kind: FutureKind,
+    ) -> hir::ExprKind<'hir> {
         let full_span = expr.span.to(await_kw_span);
 
         let is_async_gen = match self.coroutine_kind {
@@ -887,13 +898,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
         };
 
-        let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
+        let features = match await_kind {
+            FutureKind::Future => None,
+            FutureKind::AsyncIterator => Some(self.allow_for_await.clone()),
+        };
+        let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, features);
         let gen_future_span = self.mark_span_with_reason(
             DesugaringKind::Await,
             full_span,
             Some(self.allow_gen_future.clone()),
         );
-        let expr = self.lower_expr_mut(expr);
         let expr_hir_id = expr.hir_id;
 
         // Note that the name of this binding must not be changed to something else because
@@ -933,11 +947,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 hir::LangItem::GetContext,
                 arena_vec![self; task_context],
             );
-            let call = self.expr_call_lang_item_fn(
-                span,
-                hir::LangItem::FuturePoll,
-                arena_vec![self; new_unchecked, get_context],
-            );
+            let call = match await_kind {
+                FutureKind::Future => self.expr_call_lang_item_fn(
+                    span,
+                    hir::LangItem::FuturePoll,
+                    arena_vec![self; new_unchecked, get_context],
+                ),
+                FutureKind::AsyncIterator => self.expr_call_lang_item_fn(
+                    span,
+                    hir::LangItem::AsyncIteratorPollNext,
+                    arena_vec![self; new_unchecked, get_context],
+                ),
+            };
             self.arena.alloc(self.expr_unsafe(call))
         };
 
@@ -1021,11 +1042,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let awaitee_arm = self.arm(awaitee_pat, loop_expr);
 
         // `match ::std::future::IntoFuture::into_future(<expr>) { ... }`
-        let into_future_expr = self.expr_call_lang_item_fn(
-            span,
-            hir::LangItem::IntoFutureIntoFuture,
-            arena_vec![self; expr],
-        );
+        let into_future_expr = match await_kind {
+            FutureKind::Future => self.expr_call_lang_item_fn(
+                span,
+                hir::LangItem::IntoFutureIntoFuture,
+                arena_vec![self; *expr],
+            ),
+            // Not needed for `for await` because we expect to have already called
+            // `IntoAsyncIterator::into_async_iter` on it.
+            FutureKind::AsyncIterator => expr,
+        };
 
         // match <into_future_expr> {
         //     mut __awaitee => loop { .. }
@@ -1673,7 +1699,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         head: &Expr,
         body: &Block,
         opt_label: Option<Label>,
-        _loop_kind: ForLoopKind,
+        loop_kind: ForLoopKind,
     ) -> hir::Expr<'hir> {
         let head = self.lower_expr_mut(head);
         let pat = self.lower_pat(pat);
@@ -1702,17 +1728,41 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let (iter_pat, iter_pat_nid) =
             self.pat_ident_binding_mode(head_span, iter, hir::BindingAnnotation::MUT);
 
-        // `match Iterator::next(&mut iter) { ... }`
         let match_expr = {
             let iter = self.expr_ident(head_span, iter, iter_pat_nid);
-            let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
-            let next_expr = self.expr_call_lang_item_fn(
-                head_span,
-                hir::LangItem::IteratorNext,
-                arena_vec![self; ref_mut_iter],
-            );
+            let next_expr = match loop_kind {
+                ForLoopKind::For => {
+                    // `Iterator::next(&mut iter)`
+                    let ref_mut_iter = self.expr_mut_addr_of(head_span, iter);
+                    self.expr_call_lang_item_fn(
+                        head_span,
+                        hir::LangItem::IteratorNext,
+                        arena_vec![self; ref_mut_iter],
+                    )
+                }
+                ForLoopKind::ForAwait => {
+                    // we'll generate `unsafe { Pin::new_unchecked(&mut iter) })` and then pass this
+                    // to make_lowered_await with `FutureKind::AsyncIterator` which will generator
+                    // calls to `poll_next`. In user code, this would probably be a call to
+                    // `Pin::as_mut` but here it's easy enough to do `new_unchecked`.
+
+                    // `&mut iter`
+                    let iter = self.expr_mut_addr_of(head_span, iter);
+                    // `Pin::new_unchecked(...)`
+                    let iter = self.arena.alloc(self.expr_call_lang_item_fn_mut(
+                        head_span,
+                        hir::LangItem::PinNewUnchecked,
+                        arena_vec![self; iter],
+                    ));
+                    // `unsafe { ... }`
+                    let iter = self.arena.alloc(self.expr_unsafe(iter));
+                    let kind = self.make_lowered_await(head_span, iter, FutureKind::AsyncIterator);
+                    self.arena.alloc(hir::Expr { hir_id: self.next_id(), kind, span: head_span })
+                }
+            };
             let arms = arena_vec![self; none_arm, some_arm];
 
+            // `match $next_expr { ... }`
             self.expr_match(head_span, next_expr, arms, hir::MatchSource::ForLoopDesugar)
         };
         let match_stmt = self.stmt_expr(for_span, match_expr);
@@ -1732,13 +1782,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
         // `mut iter => { ... }`
         let iter_arm = self.arm(iter_pat, loop_expr);
 
-        // `match ::std::iter::IntoIterator::into_iter(<head>) { ... }`
-        let into_iter_expr = {
-            self.expr_call_lang_item_fn(
-                head_span,
-                hir::LangItem::IntoIterIntoIter,
-                arena_vec![self; head],
-            )
+        let into_iter_expr = match loop_kind {
+            ForLoopKind::For => {
+                // `::std::iter::IntoIterator::into_iter(<head>)`
+                self.expr_call_lang_item_fn(
+                    head_span,
+                    hir::LangItem::IntoIterIntoIter,
+                    arena_vec![self; head],
+                )
+            }
+            ForLoopKind::ForAwait => self.arena.alloc(head),
         };
 
         let match_expr = self.arena.alloc(self.expr_match(
@@ -2141,3 +2194,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }
     }
 }
+
+/// Used by [`LoweringContext::make_lowered_await`] to customize the desugaring based on what kind
+/// of future we are awaiting.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+enum FutureKind {
+    /// We are awaiting a normal future
+    Future,
+    /// We are awaiting something that's known to be an AsyncIterator (i.e. we are in the header of
+    /// a `for await` loop)
+    AsyncIterator,
+}
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 96ed3eee02ed9..1de5db4deb862 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -130,6 +130,7 @@ struct LoweringContext<'a, 'hir> {
     allow_try_trait: Lrc<[Symbol]>,
     allow_gen_future: Lrc<[Symbol]>,
     allow_async_iterator: Lrc<[Symbol]>,
+    allow_for_await: Lrc<[Symbol]>,
 
     /// Mapping from generics `def_id`s to TAIT generics `def_id`s.
     /// For each captured lifetime (e.g., 'a), we create a new lifetime parameter that is a generic
@@ -174,6 +175,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             } else {
                 [sym::gen_future].into()
             },
+            allow_for_await: [sym::async_iterator].into(),
             // FIXME(gen_blocks): how does `closure_track_caller`/`async_fn_track_caller`
             // interact with `gen`/`async gen` blocks
             allow_async_iterator: [sym::gen_future, sym::async_iterator].into(),
diff --git a/compiler/rustc_builtin_macros/src/assert/context.rs b/compiler/rustc_builtin_macros/src/assert/context.rs
index 4e07122a5c4c9..d244897f8a5de 100644
--- a/compiler/rustc_builtin_macros/src/assert/context.rs
+++ b/compiler/rustc_builtin_macros/src/assert/context.rs
@@ -303,7 +303,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::Continue(_)
             | ExprKind::Err
             | ExprKind::Field(_, _)
-            | ExprKind::ForLoop {..}
+            | ExprKind::ForLoop { .. }
             | ExprKind::FormatArgs(_)
             | ExprKind::IncludedBytes(..)
             | ExprKind::InlineAsm(_)
diff --git a/compiler/rustc_feature/src/unstable.rs b/compiler/rustc_feature/src/unstable.rs
index fdd4bb755e415..bb1c95c25d719 100644
--- a/compiler/rustc_feature/src/unstable.rs
+++ b/compiler/rustc_feature/src/unstable.rs
@@ -358,7 +358,7 @@ declare_features! (
     /// Allows `#[track_caller]` on async functions.
     (unstable, async_fn_track_caller, "1.73.0", Some(110011)),
     /// Allows `for await` loops.
-    (unstable, async_for_loop, "CURRENT_RUSTC_VERSION", None),
+    (unstable, async_for_loop, "CURRENT_RUSTC_VERSION", Some(118898)),
     /// Allows builtin # foo() syntax
     (unstable, builtin_syntax, "1.71.0", Some(110680)),
     /// Treat `extern "C"` function as nounwind.
diff --git a/compiler/rustc_hir/src/lang_items.rs b/compiler/rustc_hir/src/lang_items.rs
index b0b53bb7478d5..7691cd11c4f55 100644
--- a/compiler/rustc_hir/src/lang_items.rs
+++ b/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;
 
+    AsyncIteratorPollNext,   sym::async_iterator_poll_next, async_iterator_poll_next, Target::Method(MethodKind::Trait { body: false }), GenericRequirement::Exact(0);
+
     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;
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 2e7037532b5c6..95106cc64c129 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -428,6 +428,7 @@ symbols! {
         async_fn_track_caller,
         async_for_loop,
         async_iterator,
+        async_iterator_poll_next,
         atomic,
         atomic_mod,
         atomics,
@@ -894,6 +895,7 @@ symbols! {
         instruction_set,
         integer_: "integer", // underscore to avoid clashing with the function `sym::integer` below
         integral,
+        into_async_iter_into_iter,
         into_future,
         into_iter,
         intra_doc_pointers,
diff --git a/library/core/src/async_iter/async_iter.rs b/library/core/src/async_iter/async_iter.rs
index 8a45bd36f7a29..ed9cb5bfea50f 100644
--- a/library/core/src/async_iter/async_iter.rs
+++ b/library/core/src/async_iter/async_iter.rs
@@ -47,6 +47,7 @@ pub trait AsyncIterator {
     /// Rust's usual rules apply: calls must never cause undefined behavior
     /// (memory corruption, incorrect use of `unsafe` functions, or the like),
     /// regardless of the async iterator's state.
+    #[cfg_attr(not(bootstrap), lang = "async_iterator_poll_next")]
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>;
 
     /// Returns the bounds on the remaining length of the async iterator.
diff --git a/tests/ui/async-await/for-await.rs b/tests/ui/async-await/for-await.rs
new file mode 100644
index 0000000000000..6345ceb0c2798
--- /dev/null
+++ b/tests/ui/async-await/for-await.rs
@@ -0,0 +1,24 @@
+// run-pass
+// edition: 2021
+#![feature(async_iterator, async_iter_from_iter, const_waker, async_for_loop, noop_waker)]
+
+use std::future::Future;
+
+// make sure a simple for await loop works
+async fn real_main() {
+    let iter = core::async_iter::from_iter(0..3);
+    let mut count = 0;
+    for await i in iter {
+        assert_eq!(i, count);
+        count += 1;
+    }
+    assert_eq!(count, 3);
+}
+
+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) {}
+}

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 3/3] 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 f424040e4ea7e..85f785b3401be 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 6d32b3ca7007b..42247dd14b0d1 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 7e6d83fe5b6c8..38f75821772c8 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 0000000000000..c1b7c016d1fc0
--- /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 0000000000000..65bb9e88448ff
--- /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 0000000000000..48b2e8a51d830
--- /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 0000000000000..7fa133aaedcd7
--- /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) {}
+}