@@ -48,23 +48,30 @@ impl From<ast::IfExpr> for ElseBranch {
4848}
4949
5050impl ast:: IfExpr {
51- pub fn then_branch ( & self ) -> Option < ast:: BlockExpr > {
52- self . children_after_condition ( ) . next ( )
51+ pub fn condition ( & self ) -> Option < ast:: Expr > {
52+ // If the condition is a BlockExpr, check if the then body is missing.
53+ // If it is assume the condition is the expression that is missing instead.
54+ let mut exprs = support:: children ( self . syntax ( ) ) ;
55+ let first = exprs. next ( ) ;
56+ match first {
57+ Some ( ast:: Expr :: BlockExpr ( _) ) => exprs. next ( ) . and ( first) ,
58+ first => first,
59+ }
5360 }
5461
55- pub fn else_branch ( & self ) -> Option < ElseBranch > {
56- let res = match self . children_after_condition ( ) . nth ( 1 ) {
57- Some ( block) => ElseBranch :: Block ( block) ,
58- None => {
59- let elif = self . children_after_condition ( ) . next ( ) ?;
60- ElseBranch :: IfExpr ( elif)
61- }
62- } ;
63- Some ( res)
62+ pub fn then_branch ( & self ) -> Option < ast:: BlockExpr > {
63+ match support:: children ( self . syntax ( ) ) . nth ( 1 ) ? {
64+ ast:: Expr :: BlockExpr ( block) => Some ( block) ,
65+ _ => None ,
66+ }
6467 }
6568
66- fn children_after_condition < N : AstNode > ( & self ) -> impl Iterator < Item = N > {
67- self . syntax ( ) . children ( ) . skip ( 1 ) . filter_map ( N :: cast)
69+ pub fn else_branch ( & self ) -> Option < ElseBranch > {
70+ match support:: children ( self . syntax ( ) ) . nth ( 2 ) ? {
71+ ast:: Expr :: BlockExpr ( block) => Some ( ElseBranch :: Block ( block) ) ,
72+ ast:: Expr :: IfExpr ( elif) => Some ( ElseBranch :: IfExpr ( elif) ) ,
73+ _ => None ,
74+ }
6875 }
6976}
7077
0 commit comments