From e2f39dd7b95dec746c34656c0be622b9636d456e Mon Sep 17 00:00:00 2001
From: Greg Shuflin <greg.shuflin@protonmail.com>
Date: Tue, 2 Nov 2021 01:16:08 -0700
Subject: [PATCH] Fix bug with loops

---
 .../language/src/tree_walk_eval/evaluator.rs  |  9 ++++++--
 .../language/src/tree_walk_eval/test.rs       | 21 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/schala-lang/language/src/tree_walk_eval/evaluator.rs b/schala-lang/language/src/tree_walk_eval/evaluator.rs
index 172ade2..1f4bbc0 100644
--- a/schala-lang/language/src/tree_walk_eval/evaluator.rs
+++ b/schala-lang/language/src/tree_walk_eval/evaluator.rs
@@ -56,7 +56,6 @@ impl<'a, 'b> Evaluator<'a, 'b> {
     }
 
     fn block(&mut self, statements: Vec<Statement>) -> EvalResult<Primitive> {
-        //TODO need to handle breaks, returns, etc.
         let mut retval = None;
         for stmt in statements.into_iter() {
             match self.statement(stmt)? {
@@ -69,6 +68,7 @@ impl<'a, 'b> Evaluator<'a, 'b> {
                 break;
             }
             if let Some(_) = self.loop_control {
+                println!("We here?");
                 break;
             }
         }
@@ -142,6 +142,7 @@ impl<'a, 'b> Evaluator<'a, 'b> {
             Expression::Assign { ref lval, box rval } => {
                 let mem = lval.into();
                 let evaluated = self.expression(rval)?;
+                println!("Inserting {:?} into {:?}", evaluated, mem);
                 self.state.environments.insert(mem, MemoryValue::Primitive(evaluated));
                 Primitive::unit()
             }
@@ -206,6 +207,7 @@ impl<'a, 'b> Evaluator<'a, 'b> {
     ) -> EvalResult<Primitive> {
         loop {
             let cond = self.expression(cond.clone())?;
+            println!("COND: {:?}", cond);
             match cond {
                 Primitive::Literal(Literal::Bool(true)) => (),
                 Primitive::Literal(Literal::Bool(false)) => break,
@@ -214,7 +216,10 @@ impl<'a, 'b> Evaluator<'a, 'b> {
             //TODO eventually loops shoudl be able to return something
             let _output = self.block(statements.clone())?;
             match self.loop_control {
-                None | Some(LoopControlFlow::Continue) => (),
+                None => (),
+                Some(LoopControlFlow::Continue) => {
+                    self.loop_control = None;
+                }
                 Some(LoopControlFlow::Break) => {
                     break;
                 }
diff --git a/schala-lang/language/src/tree_walk_eval/test.rs b/schala-lang/language/src/tree_walk_eval/test.rs
index af7b921..ab2f3ed 100644
--- a/schala-lang/language/src/tree_walk_eval/test.rs
+++ b/schala-lang/language/src/tree_walk_eval/test.rs
@@ -467,3 +467,24 @@ count
     "#;
     eval_assert(source, "500");
 }
+
+#[test]
+fn loops_2() {
+    let source = r#"
+let mut a = 0
+let mut acc = 0
+while a < 10 {
+    acc = acc + 1
+    a = a + 1
+
+    // Without this continue, the output would be 20
+    if a == 5 then {
+        continue
+    }
+
+    acc = acc + 1
+}
+
+acc"#;
+    eval_assert(source, "19");
+}