Plumb awaitness of for loops

eholk · eholk · commit 2adf5fa1c1aa · 2023-12-12T11:50:15.000-08:00
diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
@@ -1261,7 +1261,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,
@@ -1423,10 +1423,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 }`
@@ -1529,6 +1529,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
@@ -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
@@ -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(..)
     )
@@ -46,8 +46,15 @@ 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(..) => break Some(expr),
+            Gen(..)
+            | Block(..)
+            | ForLoop { .. }
+            | If(..)
+            | Loop(..)
+            | Match(..)
+            | Struct(..)
+            | TryBlock(..)
+            | While(..) => break Some(expr),
             _ => break None,
         }
     }
diff --git 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
@@ -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);
+                // correspond to the `e.id`, so `lower_expr_for` handles attribute low}ering itself.
+                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")
                 }
 
@@ -1670,6 +1670,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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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};
@@ -1800,7 +1800,7 @@ impl<'a> Parser<'a> {
                     && matches!(
                         expr.kind,
                         ExprKind::While(_, _, None)
-                            | ExprKind::ForLoop(_, _, _, None)
+                            | ExprKind::ForLoop { label: None, .. }
                             | ExprKind::Loop(_, None, _)
                             | ExprKind::Block(_, None)
                     )
@@ -2681,8 +2681,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(..))
@@ -2695,13 +2703,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)?;
 
@@ -3783,7 +3791,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
@@ -4250,11 +4250,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
@@ -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/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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
diff --git a/src/tools/rustfmt/src/matches.rs b/src/tools/rustfmt/src/matches.rs
diff --git a/src/tools/rustfmt/src/overflow.rs b/src/tools/rustfmt/src/overflow.rs
diff --git a/src/tools/rustfmt/src/utils.rs b/src/tools/rustfmt/src/utils.rs
diff --git a/tests/ui/async-await/feature-async-for-loop.rs b/tests/ui/async-await/feature-async-for-loop.rs
diff --git a/tests/ui/async-await/feature-async-for-loop.stderr b/tests/ui/async-await/feature-async-for-loop.stderr

Original file line number	Diff line number	Diff line change
`@@ -1389,7 +1389,7 @@ pub fn noop_visit_expr<T: MutVisitor>(`
`1389`	`1389`	`vis.visit_block(body);`
`1390`	`1390`	`visit_opt(label, \|label\| vis.visit_label(label));`
`1391`	`1391`	`}`
`1392`		`- ExprKind::ForLoop(pat, iter, body, label) => {`
	`1392`	`+ ExprKind::ForLoop { pat, iter, body, label, kind: _ } => {`
`1393`	`1393`	`vis.visit_pat(pat);`
`1394`	`1394`	`vis.visit_expr(iter);`
`1395`	`1395`	`vis.visit_block(body);`
Original file line number	Diff line number	Diff line change
`@@ -852,8 +852,8 @@ trait UnusedDelimLint {`
`852`	`852`	`(cond, UnusedDelimsCtx::WhileCond, true, Some(left), Some(right), true)`
`853`	`853`	`}`
`854`	`854`
`855`		`- ForLoop(_, ref cond, ref block, ..) => {`
`856`		`- (cond, UnusedDelimsCtx::ForIterExpr, true, None, Some(block.span.lo()), true)`
	`855`	`+ ForLoop { ref iter, ref body, .. } => {`
	`856`	`+ (iter, UnusedDelimsCtx::ForIterExpr, true, None, Some(body.span.lo()), true)`
`857`	`857`	`}`
`858`	`858`
`859`	`859`	`Match(ref head, _) if Self::LINT_EXPR_IN_PATTERN_MATCHING_CTX => {`
`@@ -1085,7 +1085,7 @@ impl EarlyLintPass for UnusedParens {`
`1085`	`1085`	`}`
`1086`	`1086`
`1087`	`1087`	`match e.kind {`
`1088`		`- ExprKind::Let(ref pat, _, _, _) \| ExprKind::ForLoop(ref pat, ..) => {`
	`1088`	`+ ExprKind::Let(ref pat, _, _, _) \| ExprKind::ForLoop { ref pat, .. } => {`
`1089`	`1089`	`self.check_unused_parens_pat(cx, pat, false, false, (true, true));`
`1090`	`1090`	`}`
`1091`	`1091`	// We ignore parens in cases like `if (((let Some(0) = Some(1))))` because we already