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

---
 src/closures.rs |  4 ++--
 src/expr.rs     | 24 +++++++++++++++++-------
 src/matches.rs  |  2 +-
 src/overflow.rs |  2 +-
 src/utils.rs    |  4 ++--
 5 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/src/closures.rs b/src/closures.rs
index f698f494ae5..5bf29441b54 100644
--- a/src/closures.rs
+++ b/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/expr.rs b/src/expr.rs
index a68bd6694ba..7808f891336 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -3,7 +3,7 @@
 
 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 @@ fn new_for(
         block: &'a ast::Block,
         label: Option<ast::Label>,
         span: Span,
+        kind: ForLoopKind,
     ) -> ControlFlow<'a> {
         ControlFlow {
             cond: Some(cond),
@@ -778,7 +785,10 @@ fn new_for(
             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/matches.rs b/src/matches.rs
index ef509b56837..5a00984d4c0 100644
--- a/src/matches.rs
+++ b/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/overflow.rs b/src/overflow.rs
index d81bf24dbd1..f46583b1c89 100644
--- a/src/overflow.rs
+++ b/src/overflow.rs
@@ -409,7 +409,7 @@ fn rewrite_last_item_with_overflow(
                     // 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/utils.rs b/src/utils.rs
index 7d7bbf11529..642b6603b1e 100644
--- a/src/utils.rs
+++ b/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)