From 87c4a4104da2acc662f777324f4ae3c6702bb5cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lauren=C8=9Biu=20Nicola?= <lnicola@dend.ro>
Date: Sat, 26 Feb 2022 16:45:06 +0200
Subject: [PATCH] Fix body selection in while loops

---
 crates/hir_ty/src/tests/regression.rs    | 22 ++++++++++++++++++++++
 crates/syntax/src/ast/generated/nodes.rs |  1 -
 crates/syntax/src/ast/node_ext.rs        |  9 +++++++++
 crates/syntax/src/tests/sourcegen_ast.rs |  5 +++--
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/crates/hir_ty/src/tests/regression.rs b/crates/hir_ty/src/tests/regression.rs
index 4a0a5d6fa1..812f12aced 100644
--- a/crates/hir_ty/src/tests/regression.rs
+++ b/crates/hir_ty/src/tests/regression.rs
@@ -1245,3 +1245,25 @@ fn test() {
         "#]],
     );
 }
+
+#[test]
+fn while_loop_block_expr_iterable() {
+    check_infer(
+        r#"
+fn test() {
+    while { true } {
+        let y = 0;
+    }
+}
+        "#,
+        expect![[r#"
+            10..59 '{     ...   } }': ()
+            16..57 'while ...     }': ()
+            22..30 '{ true }': bool
+            24..28 'true': bool
+            31..57 '{     ...     }': ()
+            45..46 'y': i32
+            49..50 '0': i32
+        "#]],
+    );
+}
diff --git a/crates/syntax/src/ast/generated/nodes.rs b/crates/syntax/src/ast/generated/nodes.rs
index 6c4729ef36..11247fbf1c 100644
--- a/crates/syntax/src/ast/generated/nodes.rs
+++ b/crates/syntax/src/ast/generated/nodes.rs
@@ -1035,7 +1035,6 @@ pub struct WhileExpr {
     pub(crate) syntax: SyntaxNode,
 }
 impl ast::HasAttrs for WhileExpr {}
-impl ast::HasLoopBody for WhileExpr {}
 impl WhileExpr {
     pub fn while_token(&self) -> Option<SyntaxToken> { support::token(&self.syntax, T![while]) }
     pub fn condition(&self) -> Option<Expr> { support::child(&self.syntax) }
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index 1b916e91bf..97455abd80 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -769,6 +769,15 @@ impl ast::HasLoopBody for ast::ForExpr {
     }
 }
 
+impl ast::HasLoopBody for ast::WhileExpr {
+    fn loop_body(&self) -> Option<ast::BlockExpr> {
+        let mut exprs = support::children(self.syntax());
+        let first = exprs.next();
+        let second = exprs.next();
+        second.or(first)
+    }
+}
+
 impl ast::HasAttrs for ast::AnyHasDocComments {}
 
 impl From<ast::Adt> for ast::Item {
diff --git a/crates/syntax/src/tests/sourcegen_ast.rs b/crates/syntax/src/tests/sourcegen_ast.rs
index 4a608da9ff..5ed56a81c4 100644
--- a/crates/syntax/src/tests/sourcegen_ast.rs
+++ b/crates/syntax/src/tests/sourcegen_ast.rs
@@ -85,8 +85,9 @@ fn generate_nodes(kinds: KindsSrc<'_>, grammar: &AstSrc) -> String {
                 .traits
                 .iter()
                 .filter(|trait_name| {
-                    // For loops have two expressions so this might collide, therefor manual impl it
-                    node.name != "ForExpr" || trait_name.as_str() != "HasLoopBody"
+                    // Loops have two expressions so this might collide, therefor manual impl it
+                    node.name != "ForExpr" && node.name != "WhileExpr"
+                        || trait_name.as_str() != "HasLoopBody"
                 })
                 .map(|trait_name| {
                     let trait_name = format_ident!("{}", trait_name);