:arrow_up: rust-analyzer

2 files changed, 29 insertions, 20 deletions

diff --git a/crates/syntax/src/ast/expr_ext.rs b/crates/syntax/src/ast/expr_ext.rs
index db66d08a73b..c43d0830b9e 100644
--- a/crates/syntax/src/ast/expr_ext.rs
+++ b/crates/syntax/src/ast/expr_ext.rs
@@ -48,23 +48,30 @@ impl From<ast::IfExpr> for ElseBranch {
 }
 
 impl ast::IfExpr {
-    pub fn then_branch(&self) -> Option<ast::BlockExpr> {
-        self.children_after_condition().next()
+    pub fn condition(&self) -> Option<ast::Expr> {
+        // If the condition is a BlockExpr, check if the then body is missing.
+        // If it is assume the condition is the expression that is missing instead.
+        let mut exprs = support::children(self.syntax());
+        let first = exprs.next();
+        match first {
+            Some(ast::Expr::BlockExpr(_)) => exprs.next().and(first),
+            first => first,
+        }
     }
 
-    pub fn else_branch(&self) -> Option<ElseBranch> {
-        let res = match self.children_after_condition().nth(1) {
-            Some(block) => ElseBranch::Block(block),
-            None => {
-                let elif = self.children_after_condition().next()?;
-                ElseBranch::IfExpr(elif)
-            }
-        };
-        Some(res)
+    pub fn then_branch(&self) -> Option<ast::BlockExpr> {
+        match support::children(self.syntax()).nth(1)? {
+            ast::Expr::BlockExpr(block) => Some(block),
+            _ => None,
+        }
     }
 
-    fn children_after_condition<N: AstNode>(&self) -> impl Iterator<Item = N> {
-        self.syntax().children().skip(1).filter_map(N::cast)
+    pub fn else_branch(&self) -> Option<ElseBranch> {
+        match support::children(self.syntax()).nth(2)? {
+            ast::Expr::BlockExpr(block) => Some(ElseBranch::Block(block)),
+            ast::Expr::IfExpr(elif) => Some(ElseBranch::IfExpr(elif)),
+            _ => None,
+        }
     }
 }
 
@@ -356,7 +363,15 @@ impl ast::BlockExpr {
             Some(it) => it,
             None => return true,
         };
-        !matches!(parent.kind(), FN | IF_EXPR | WHILE_EXPR | LOOP_EXPR)
+        match parent.kind() {
+            FOR_EXPR | IF_EXPR => parent
+                .children()
+                .filter(|it| ast::Expr::can_cast(it.kind()))
+                .next()
+                .map_or(true, |it| it == *self.syntax()),
+            LET_ELSE | FN | WHILE_EXPR | LOOP_EXPR | CONST_BLOCK_PAT => false,
+            _ => true,
+        }
     }
 }
 
diff --git a/crates/syntax/src/ast/node_ext.rs b/crates/syntax/src/ast/node_ext.rs
index 15bd5ab3c72..3308077da5b 100644
--- a/crates/syntax/src/ast/node_ext.rs
+++ b/crates/syntax/src/ast/node_ext.rs
@@ -937,12 +937,6 @@ impl From<ast::Adt> for ast::Item {
     }
 }
 
-impl ast::IfExpr {
-    pub fn condition(&self) -> Option<ast::Expr> {
-        support::child(&self.syntax)
-    }
-}
-
 impl ast::MatchGuard {
     pub fn condition(&self) -> Option<ast::Expr> {
         support::child(&self.syntax)