From 9fe206956f791a7974d3ba3a2534c1e32b09a3a0 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Wed, 15 Mar 2023 12:53:39 +0100
Subject: [PATCH] fix: Fix ast::IfExpr child accessors

---
 crates/hir-expand/src/fixup.rs    |  3 +--
 crates/syntax/src/ast/expr_ext.rs | 33 +++++++++++++++++++------------
 crates/syntax/src/ast/node_ext.rs |  6 ------
 3 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/crates/hir-expand/src/fixup.rs b/crates/hir-expand/src/fixup.rs
index c811d1c66a8..b273f21768c 100644
--- a/crates/hir-expand/src/fixup.rs
+++ b/crates/hir-expand/src/fixup.rs
@@ -636,9 +636,8 @@ fn foo() {
     if {}
 }
 "#,
-            // the {} gets parsed as the condition, I think?
             expect![[r#"
-fn foo () {if {} {}}
+fn foo () {if __ra_fixup {} {}}
 "#]],
         )
     }
diff --git a/crates/syntax/src/ast/expr_ext.rs b/crates/syntax/src/ast/expr_ext.rs
index 699c68aeece..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 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 then_branch(&self) -> Option<ast::BlockExpr> {
-        self.children_after_condition().next()
+        match support::children(self.syntax()).nth(1)? {
+            ast::Expr::BlockExpr(block) => Some(block),
+            _ => None,
+        }
     }
 
     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)
-    }
-
-    fn children_after_condition<N: AstNode>(&self) -> impl Iterator<Item = N> {
-        self.syntax().children().skip(1).filter_map(N::cast)
+        match support::children(self.syntax()).nth(2)? {
+            ast::Expr::BlockExpr(block) => Some(ElseBranch::Block(block)),
+            ast::Expr::IfExpr(elif) => Some(ElseBranch::IfExpr(elif)),
+            _ => None,
+        }
     }
 }
 
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)