Unnamed repository; edit this file 'description' to name the repository.
Diffstat (limited to 'crates/syntax/src/ast/expr_ext.rs')
| -rw-r--r-- | crates/syntax/src/ast/expr_ext.rs | 33 |
1 files changed, 20 insertions, 13 deletions
diff --git a/crates/syntax/src/ast/expr_ext.rs b/crates/syntax/src/ast/expr_ext.rs index 699c68aeec..c43d0830b9 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, + } } } |