From 2beb493aadffe9dbb1bd10eff23670d6ef5bd8fe Mon Sep 17 00:00:00 2001
From: mars <mars@tebibyte.media>
Date: Wed, 2 Mar 2022 13:20:27 -0700
Subject: [PATCH] Support if statements AND expressions (pogchamp)

---
 src/parse.rs | 144 +++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 106 insertions(+), 38 deletions(-)

diff --git a/src/parse.rs b/src/parse.rs
index bd54d21..19d92f2 100644
--- a/src/parse.rs
+++ b/src/parse.rs
@@ -283,31 +283,49 @@ impl<'a> BranchBody<'a> {
     pub fn build(lexer: &mut Lexer<'a>) -> Self {
         let mut statements = Vec::new();
 
-        let mut tail_expr = None;
-        loop {
-            let tok = lexer.next().unwrap();
+        let mut next = None;
+        let tail_expr = loop {
+            let tok = if let Some(tok) = next {
+                next = None;
+                tok
+            } else {
+                lexer.next().unwrap()
+            };
+
             match tok {
                 Token::Let => {
                     let (statement, next) = Statement::build_let(lexer);
                     statements.push(statement);
                     match next {
                         Token::Semicolon => {}
-                        Token::BraceClose => break,
+                        Token::BraceClose => break None,
                         _ => lexer.panic_message("Unexpected token"),
                     }
                 }
-                Token::BraceClose => break,
+                Token::If => {
+                    let (stmt, tail) = IfStmt::build(lexer);
+                    if tail == Token::BraceClose {
+                        if let Some(_) = stmt.else_body {
+                            let expr = stmt.to_expr();
+                            break Some(Expr::If(expr));
+                        } else {
+                            statements.push(Statement::If(stmt));
+                            break None;
+                        }
+                    } else {
+                        statements.push(Statement::If(stmt));
+                        next = Some(tail);
+                    }
+                }
+                Token::BraceClose => break None,
                 _ => match Expr::build_start(tok, lexer) {
                     (None, Token::Semicolon | Token::BraceClose) => {}
                     (Some(expr), Token::Semicolon) => statements.push(Statement::Expr(expr)),
-                    (Some(expr), Token::BraceClose) => {
-                        tail_expr = Some(expr);
-                        break;
-                    }
+                    (Some(expr), Token::BraceClose) => break Some(expr),
                     _ => lexer.panic_message("Unexpected token"),
                 },
             }
-        }
+        };
 
         Self {
             statements,
@@ -327,11 +345,7 @@ pub enum Expr<'a> {
     SelfMember(&'a str),
     Group(Box<Expr<'a>>),
     FnCall(Box<Expr<'a>>, Vec<Expr<'a>>),
-    If {
-        test_expr: Box<Expr<'a>>,
-        then_body: BranchBody<'a>,
-        else_body: BranchBody<'a>,
-    },
+    If(IfExpr<'a>),
 }
 
 impl<'a> Expr<'a> {
@@ -346,24 +360,7 @@ impl<'a> Expr<'a> {
                 }
                 Self::Group(Box::new(inner.unwrap()))
             }
-            Token::If => {
-                let test_expr = match Self::build(lexer) {
-                    (Some(test_expr), Token::BraceOpen) => Box::new(test_expr),
-                    (None, Token::BraceOpen) => lexer.panic_message("Expected test expression"),
-                    _ => lexer.panic_message("Expected opening brace"),
-                };
-
-                let then_body = BranchBody::build(lexer);
-                lexer.eat_expect(Token::Else);
-                lexer.eat_expect(Token::BraceOpen);
-                let else_body = BranchBody::build(lexer);
-
-                Self::If {
-                    test_expr,
-                    then_body,
-                    else_body,
-                }
-            }
+            Token::If => Self::If(IfExpr::build(lexer)),
             Token::BinaryInteger => Self::Literal(Literal::BinaryInteger(lexer.slice())),
             Token::OctalInteger => Self::Literal(Literal::OctalInteger(lexer.slice())),
             Token::HexInteger => Self::Literal(Literal::HexInteger(lexer.slice())),
@@ -485,16 +482,12 @@ pub enum Literal<'a> {
 #[derive(Debug)]
 pub enum Statement<'a> {
     Expr(Expr<'a>),
+    If(IfStmt<'a>),
     Let {
         var: &'a str,
         mutable: bool,
         expr: Expr<'a>,
     },
-    If {
-        test_expr: Expr<'a>,
-        then_body: BranchBody<'a>,
-        else_body: Option<BranchBody<'a>>,
-    },
 }
 
 impl<'a> Statement<'a> {
@@ -520,6 +513,81 @@ impl<'a> Statement<'a> {
     }
 }
 
+#[derive(Debug)]
+pub struct IfExpr<'a> {
+    pub test_expr: Box<Expr<'a>>,
+    pub then_body: BranchBody<'a>,
+    pub else_body: BranchBody<'a>,
+}
+
+impl<'a> IfExpr<'a> {
+    pub fn build(lexer: &mut Lexer<'a>) -> Self {
+        let test_expr = Box::new(Self::eat_test_expr(lexer));
+        let then_body = BranchBody::build(lexer);
+        lexer.eat_expect(Token::Else);
+        lexer.eat_expect(Token::BraceOpen);
+        let else_body = BranchBody::build(lexer);
+        Self {
+            test_expr,
+            then_body,
+            else_body,
+        }
+    }
+
+    pub fn eat_test_expr(lexer: &mut Lexer<'a>) -> Expr<'a> {
+        match Expr::build(lexer) {
+            (Some(test_expr), Token::BraceOpen) => test_expr,
+            (Some(_), _) => lexer.panic_message("Expected opening brace"),
+            _ => lexer.panic_message("Expected test expression"),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct IfStmt<'a> {
+    pub test_expr: Expr<'a>,
+    pub then_body: BranchBody<'a>,
+    pub else_body: Option<BranchBody<'a>>,
+}
+
+impl<'a> IfStmt<'a> {
+    pub fn build(lexer: &mut Lexer<'a>) -> (Self, Token) {
+        let test_expr = IfExpr::eat_test_expr(lexer);
+        let then_body = BranchBody::build(lexer);
+
+        let (else_body, tail) = match lexer.next().unwrap() {
+            Token::Else => {
+                lexer.eat_expect(Token::BraceOpen);
+                (Some(BranchBody::build(lexer)), lexer.next().unwrap())
+            }
+            next => (None, next),
+        };
+
+        (
+            Self {
+                test_expr,
+                then_body,
+                else_body,
+            },
+            tail,
+        )
+    }
+
+    pub fn to_expr(self) -> IfExpr<'a> {
+        let Self {
+            test_expr,
+            then_body,
+            else_body,
+        } = self;
+
+        IfExpr {
+            test_expr: Box::new(test_expr),
+            then_body,
+            else_body: else_body.unwrap(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;