fix AST for if expressions

then is not always a block...

2 files changed, 60 insertions, 2 deletions

diff --git a/crates/ra_syntax/src/ast.rs b/crates/ra_syntax/src/ast.rs
index 00c60ebf3..ab3dd1b84 100644
--- a/crates/ra_syntax/src/ast.rs
+++ b/crates/ra_syntax/src/ast.rs
@@ -285,13 +285,27 @@ impl LetStmt {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ElseBranchFlavor<'a> {
+    Block(&'a Block),
+    IfExpr(&'a IfExpr),
+}
+
 impl IfExpr {
     pub fn then_branch(&self) -> Option<&Block> {
         self.blocks().nth(0)
     }
-    pub fn else_branch(&self) -> Option<&Block> {
-        self.blocks().nth(1)
+    pub fn else_branch(&self) -> Option<ElseBranchFlavor> {
+        let res = match self.blocks().nth(1) {
+            Some(block) => ElseBranchFlavor::Block(block),
+            None => {
+                let elif: &IfExpr = child_opt(self)?;
+                ElseBranchFlavor::IfExpr(elif)
+            }
+        };
+        Some(res)
     }
+
     fn blocks(&self) -> AstChildren<Block> {
         children(self)
     }
diff --git a/crates/ra_syntax/src/ast/generated.rs b/crates/ra_syntax/src/ast/generated.rs
index 3ace6533c..4f8723ae7 100644
--- a/crates/ra_syntax/src/ast/generated.rs
+++ b/crates/ra_syntax/src/ast/generated.rs
@@ -660,6 +660,50 @@ impl ToOwned for DynTraitType {
 
 impl DynTraitType {}
 
+// ElseBranch
+#[derive(Debug, PartialEq, Eq, Hash)]
+#[repr(transparent)]
+pub struct ElseBranch {
+    pub(crate) syntax: SyntaxNode,
+}
+unsafe impl TransparentNewType for ElseBranch {
+    type Repr = rowan::SyntaxNode<RaTypes>;
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ElseBranchKind<'a> {
+    Block(&'a Block),
+    IfExpr(&'a IfExpr),
+}
+
+impl AstNode for ElseBranch {
+    fn cast(syntax: &SyntaxNode) -> Option<&Self> {
+        match syntax.kind() {
+            | BLOCK
+            | IF_EXPR => Some(ElseBranch::from_repr(syntax.into_repr())),
+            _ => None,
+        }
+    }
+    fn syntax(&self) -> &SyntaxNode { &self.syntax }
+}
+
+impl ToOwned for ElseBranch {
+    type Owned = TreeArc<ElseBranch>;
+    fn to_owned(&self) -> TreeArc<ElseBranch> { TreeArc::cast(self.syntax.to_owned()) }
+}
+
+impl ElseBranch {
+    pub fn kind(&self) -> ElseBranchKind {
+        match self.syntax.kind() {
+            BLOCK => ElseBranchKind::Block(Block::cast(&self.syntax).unwrap()),
+            IF_EXPR => ElseBranchKind::IfExpr(IfExpr::cast(&self.syntax).unwrap()),
+            _ => unreachable!(),
+        }
+    }
+}
+
+impl ElseBranch {}
+
 // EnumDef
 #[derive(Debug, PartialEq, Eq, Hash)]
 #[repr(transparent)]