From b7288ceed8d7ec901e89022b0593a3752ce51814 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Fri, 8 Mar 2024 14:29:53 -0500
Subject: [PATCH] Implement if/else, loops and breaks

---
 src/abstract_tree/if_else.rs   |  8 ++++++--
 src/abstract_tree/loop.rs      | 11 ++++++-----
 src/abstract_tree/statement.rs | 16 ++++++++++------
 src/parser.rs                  | 23 +++++++++++++++++++----
 tests/statements.rs            | 16 ++++++++++++++++
 5 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/src/abstract_tree/if_else.rs b/src/abstract_tree/if_else.rs
index 5a4ac10..e42f03c 100644
--- a/src/abstract_tree/if_else.rs
+++ b/src/abstract_tree/if_else.rs
@@ -28,11 +28,15 @@ impl<'src> IfElse<'src> {
 
 impl<'src> AbstractTree for IfElse<'src> {
     fn expected_type(&self, _context: &Context) -> Result<Type, ValidationError> {
-        todo!()
+        Ok(Type::None)
     }
 
     fn validate(&self, _context: &Context) -> Result<(), ValidationError> {
-        todo!()
+        if let Type::Boolean = self.if_expression.expected_type(_context)? {
+            Ok(())
+        } else {
+            Err(ValidationError::ExpectedBoolean)
+        }
     }
 
     fn run(self, _context: &Context) -> Result<Action, RuntimeError> {
diff --git a/src/abstract_tree/loop.rs b/src/abstract_tree/loop.rs
index 2767cd6..3be2eca 100644
--- a/src/abstract_tree/loop.rs
+++ b/src/abstract_tree/loop.rs
@@ -40,11 +40,12 @@ impl<'src> AbstractTree for Loop<'src> {
             }
 
             let statement = self.statements[index].clone();
+            let action = statement.run(_context)?;
 
-            if let Statement::Break(expression) = statement {
-                break expression.run(_context);
-            } else {
-                statement.run(_context)?;
+            match action {
+                Action::Return(_) => {}
+                Action::None => {}
+                r#break => return Ok(r#break),
             }
         }
     }
@@ -66,6 +67,6 @@ mod tests {
         }
         .run(&Context::new());
 
-        assert_eq!(result, Ok(Action::Return(Value::integer(42))))
+        assert_eq!(result, Ok(Action::Break(Value::integer(42))))
     }
 }
diff --git a/src/abstract_tree/statement.rs b/src/abstract_tree/statement.rs
index 5791771..e7d3c86 100644
--- a/src/abstract_tree/statement.rs
+++ b/src/abstract_tree/statement.rs
@@ -22,7 +22,7 @@ impl<'src> AbstractTree for Statement<'src> {
             Statement::Block(block) => block.expected_type(_context),
             Statement::Break(expression) => expression.expected_type(_context),
             Statement::Expression(expression) => expression.expected_type(_context),
-            Statement::IfElse(_) => todo!(),
+            Statement::IfElse(if_else) => if_else.expected_type(_context),
             Statement::Loop(r#loop) => r#loop.expected_type(_context),
         }
     }
@@ -30,10 +30,10 @@ impl<'src> AbstractTree for Statement<'src> {
     fn validate(&self, _context: &Context) -> Result<(), ValidationError> {
         match self {
             Statement::Assignment(assignment) => assignment.validate(_context),
-            Statement::Block(_) => todo!(),
+            Statement::Block(block) => block.validate(_context),
             Statement::Break(expression) => expression.validate(_context),
             Statement::Expression(expression) => expression.validate(_context),
-            Statement::IfElse(_) => todo!(),
+            Statement::IfElse(if_else) => if_else.validate(_context),
             Statement::Loop(r#loop) => r#loop.validate(_context),
         }
     }
@@ -41,10 +41,14 @@ impl<'src> AbstractTree for Statement<'src> {
     fn run(self, _context: &Context) -> Result<Action, RuntimeError> {
         match self {
             Statement::Assignment(assignment) => assignment.run(_context),
-            Statement::Block(_) => todo!(),
-            Statement::Break(expression) => expression.run(_context),
+            Statement::Block(block) => block.run(_context),
+            Statement::Break(expression) => {
+                let value = expression.run(_context)?.as_return_value()?;
+
+                Ok(Action::Break(value))
+            }
             Statement::Expression(expression) => expression.run(_context),
-            Statement::IfElse(_) => todo!(),
+            Statement::IfElse(if_else) => if_else.run(_context),
             Statement::Loop(r#loop) => r#loop.run(_context),
         }
     }
diff --git a/src/parser.rs b/src/parser.rs
index a56a930..b7b2738 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -259,10 +259,13 @@ pub fn parser<'src>() -> DustParser<'src> {
         let if_else = just(Token::Keyword("if"))
             .ignore_then(expression.clone())
             .then(statement.clone())
-            .then_ignore(just(Token::Keyword("else")))
-            .then(statement.clone().or_not())
-            .map(|((if_expression, if_block), else_block)| {
-                Statement::IfElse(IfElse::new(if_expression, if_block, else_block))
+            .then(
+                just(Token::Keyword("else"))
+                    .ignore_then(statement.clone())
+                    .or_not(),
+            )
+            .map(|((if_expression, if_statement), else_statement)| {
+                Statement::IfElse(IfElse::new(if_expression, if_statement, else_statement))
             })
             .boxed();
 
@@ -290,6 +293,18 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn r#if() {
+        assert_eq!(
+            parse(&lex("if true 'foo'").unwrap()).unwrap()[0].0,
+            Statement::IfElse(IfElse::new(
+                Expression::Value(ValueNode::Boolean(true)),
+                Statement::Expression(Expression::Value(ValueNode::String("foo"))),
+                None
+            ))
+        )
+    }
+
     #[test]
     fn if_else() {
         assert_eq!(
diff --git a/tests/statements.rs b/tests/statements.rs
index a9ce096..b0e3caa 100644
--- a/tests/statements.rs
+++ b/tests/statements.rs
@@ -18,3 +18,19 @@ fn loops_and_breaks() {
         Ok(Some(Value::string("foobar")))
     )
 }
+
+#[test]
+fn r#if() {
+    assert_eq!(
+        interpret("if true 'foobar'"),
+        Ok(Some(Value::string("foobar")))
+    )
+}
+
+#[test]
+fn if_else() {
+    assert_eq!(
+        interpret("if false 'foo' else 'bar'"),
+        Ok(Some(Value::string("bar")))
+    )
+}