From 28c65b07156ce805f24687f7e7f667a50cdbef3b Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sun, 11 Aug 2024 15:03:26 -0400
Subject: [PATCH] Clean up control flow code in the VM; Remove returns from if
 and if/else_if statements

---
 dust-lang/src/parse.rs |  2 --
 dust-lang/src/vm.rs    | 61 +++++++++++++++++++++++++++++++-----------
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/dust-lang/src/parse.rs b/dust-lang/src/parse.rs
index ce32261..9ad37c3 100644
--- a/dust-lang/src/parse.rs
+++ b/dust-lang/src/parse.rs
@@ -226,8 +226,6 @@ impl<'src> Parser<'src> {
                             if let Token::Else = self.current.0 {
                                 self.next_token()?;
                             } else {
-                                let else_if_end = self.current.1;
-
                                 return Ok(Node::new(
                                     Statement::IfElseIf {
                                         condition,
diff --git a/dust-lang/src/vm.rs b/dust-lang/src/vm.rs
index a20699c..6bb166c 100644
--- a/dust-lang/src/vm.rs
+++ b/dust-lang/src/vm.rs
@@ -254,15 +254,16 @@ impl Vm {
                         position: condition_position,
                     });
                 };
-
-                if let Some(condition) = condition_value.as_boolean() {
-                    if condition {
-                        return self.run_node(*body, context);
-                    }
+                let condition = if let Some(condition) = condition_value.as_boolean() {
+                    condition
                 } else {
                     return Err(VmError::ExpectedBoolean {
                         position: condition_position,
                     });
+                };
+
+                if condition {
+                    self.run_node(*body, context)?;
                 }
 
                 Ok(None)
@@ -321,15 +322,16 @@ impl Vm {
                                         position: condition_position,
                                     });
                                 };
-
-                            if let Some(condition) = condition_value.as_boolean() {
-                                if condition {
-                                    return self.run_node(body, context);
-                                }
+                            let condition = if let Some(condition) = condition_value.as_boolean() {
+                                condition
                             } else {
                                 return Err(VmError::ExpectedBoolean {
                                     position: condition_position,
                                 });
+                            };
+
+                            if condition {
+                                self.run_node(body, context)?;
                             }
                         }
 
@@ -370,15 +372,16 @@ impl Vm {
                                         position: condition_position,
                                     });
                                 };
-
-                            if let Some(condition) = condition_value.as_boolean() {
-                                if condition {
-                                    return self.run_node(body, context);
-                                }
+                            let condition = if let Some(condition) = condition_value.as_boolean() {
+                                condition
                             } else {
                                 return Err(VmError::ExpectedBoolean {
                                     position: condition_position,
                                 });
+                            };
+
+                            if condition {
+                                return self.run_node(body, context);
                             }
                         }
 
@@ -657,6 +660,34 @@ impl Display for VmError {
 mod tests {
     use super::*;
 
+    #[test]
+    fn r#if() {
+        let input = "if true { 1 }";
+
+        assert_eq!(run(input, &mut Context::new()), Ok(None));
+    }
+
+    #[test]
+    fn if_else() {
+        let input = "if false { 1 } else { 2 }";
+
+        assert_eq!(run(input, &mut Context::new()), Ok(Some(Value::integer(2))));
+    }
+
+    #[test]
+    fn if_else_if() {
+        let input = "if false { 1 } else if true { 2 }";
+
+        assert_eq!(run(input, &mut Context::new()), Ok(None));
+    }
+
+    #[test]
+    fn if_else_if_else() {
+        let input = "if false { 1 } else if false { 2 } else { 3 }";
+
+        assert_eq!(run(input, &mut Context::new()), Ok(Some(Value::integer(3))));
+    }
+
     #[test]
     fn while_loop() {
         let input = "x = 0; while x < 5 { x += 1; } x";