Refactor if statement, introduce nodes for elif and else

Author: konstin

ast/src/generic.rs

@@ -606,12 +606,12 @@ pub struct StmtIf {
     pub range: TextRange,
     pub test: Box<Expr>,
     pub body: Vec<Stmt>,
-    pub orelse: Vec<Stmt>,
+    pub elif_else_clauses: Vec<ElifElseClause>,
 }
 
 impl Node for StmtIf {
     const NAME: &'static str = "If";
-    const FIELD_NAMES: &'static [&'static str] = &["test", "body", "orelse"];
+    const FIELD_NAMES: &'static [&'static str] = &["test", "body", "elif_else_clauses"];
 }
 impl From<StmtIf> for Stmt {
     fn from(payload: StmtIf) -> Self {
@@ -624,6 +624,18 @@ impl From<StmtIf> for Ast {
     }
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct ElifElseClause {
+    pub range: TextRange,
+    pub test: Option<Expr>,
+    pub body: Vec<Stmt>,
+}
+
+impl Node for ElifElseClause {
+    const NAME: &'static str = "ElifElse";
+    const FIELD_NAMES: &'static [&'static str] = &["test", "body"];
+}
+
 /// See also [With](https://docs.python.org/3/library/ast.html#ast.With)
 #[derive(Clone, Debug, PartialEq)]
 pub struct StmtWith {

ast/src/ranged.rs

@@ -122,6 +122,11 @@ impl Ranged for crate::generic::StmtIf {
         self.range
     }
 }
+impl Ranged for crate::generic::ElifElseClause {
+    fn range(&self) -> TextRange {
+        self.range
+    }
+}
 impl Ranged for crate::generic::StmtWith {
     fn range(&self) -> TextRange {
         self.range

parser/src/python.lalrpop

@@ -759,25 +759,24 @@ ClassPattern: ast::Pattern = {
 }
 
 IfStatement: ast::Stmt = {
-    <location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(<@L> "elif" <NamedExpressionTest> ":" <Suite>)*> <s3:("else" ":" <Suite>)?> => {
-        // Determine last else:
-        let mut last = s3.unwrap_or_default();
-        let end_location = last
+    <location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(<@L> "elif" <NamedExpressionTest> ":" <Suite>)*> <s3:(<@L> "else" ":" <Suite>)?> => {
+        let elif_else_clauses: Vec<_> = s2.into_iter().map(|(start, test, body)| ast::ElifElseClause {
+            range: (start..body.last().unwrap().end()).into(),
+            test: Some(test),
+            body,
+        }).chain(s3.into_iter().map(|(start, body)| ast::ElifElseClause {
+            range: (start..body.last().unwrap().end()).into(),
+            test: None,
+            body,
+        })).collect();
+
+        let end_location = elif_else_clauses
             .last()
-            .or_else(|| s2.last().and_then(|last| last.2.last()))
-            .or_else(|| body.last())
-            .unwrap()
-            .end();
-        // handle elif:
-        for i in s2.into_iter().rev() {
-            let x = ast::Stmt::If(
-                ast::StmtIf { test: Box::new(i.1), body: i.2, orelse: last, range: (i.0..end_location).into() }
-            );
-            last = vec![x];
-        }
+            .map(|last| last.end())
+            .unwrap_or(body.last().unwrap().end());
 
         ast::Stmt::If(
-            ast::StmtIf { test: Box::new(test), body, orelse: last, range: (location..end_location).into() }
+            ast::StmtIf { test: Box::new(test), body, elif_else_clauses, range: (location..end_location).into() }
         )
     },
 };