@@ -48,23 +48,30 @@ impl From<ast::IfExpr> for ElseBranch {
48
48
}
49
49
50
50
impl 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
+ }
53
60
}
54
61
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
+ }
64
67
}
65
68
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
+ }
68
75
}
69
76
}
70
77
0 commit comments