Merge pull request #60 from FredTheDino/break-continue

break continue

4 files changed, 168 insertions, 13 deletions

diff --git a/src/compiler.rs b/src/compiler.rs
index eb6058a..f2daa6e 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -129,14 +129,72 @@ struct Variable {
     mutable: bool,
 }
 
+enum LoopOp {
+    Continue,
+    Break,
+}
+
 struct Frame {
+    loops: Vec<Vec<(usize, usize, LoopOp)>>,
     stack: Vec<Variable>,
     upvalues: Vec<Variable>,
     scope: usize,
-    variables_below: usize,
 }
 
 impl Frame {
+    fn new() -> Self {
+        Self {
+            loops: Vec::new(),
+            stack: Vec::new(),
+            upvalues: Vec::new(),
+            scope: 0,
+        }
+    }
+
+    fn count_this_scope(&self) -> usize {
+        for (i, var) in self.stack.iter().rev().enumerate() {
+            println!("i:{} - {} == {}", i, var.scope, self.scope);
+            if var.scope != self.scope {
+                // return i;
+            }
+        }
+        return self.stack.len();
+    }
+
+    fn push_loop(&mut self) {
+        self.loops.push(Vec::new());
+    }
+
+    fn pop_loop(&mut self, block: &mut Block, stacktarget: usize, start: usize, end: usize) {
+        // Compiler error if this fails
+        for (addr, stacksize, op) in self.loops.pop().unwrap().iter() {
+            let to_pop = stacksize - stacktarget;
+            let op = match op {
+                LoopOp::Continue => Op::JmpNPop(start, to_pop),
+                LoopOp::Break => Op::JmpNPop(end, to_pop),
+            };
+            block.patch(op, *addr);
+        }
+    }
+
+    fn add_continue(&mut self, addr: usize, stacksize: usize, block: &mut Block) -> Result<(), ()> {
+        if let Some(top) = self.loops.last_mut() {
+            top.push((addr, stacksize, LoopOp::Continue));
+            Ok(())
+        } else {
+            Err(())
+        }
+    }
+
+    fn add_break(&mut self, addr: usize, stacksize: usize, block: &mut Block) -> Result<(), ()> {
+        if let Some(top) = self.loops.last_mut() {
+            top.push((addr, stacksize, LoopOp::Break));
+            Ok(())
+        } else {
+            Err(())
+        }
+    }
+
     fn find_local(&self, name: &str) -> Option<Variable> {
         for var in self.stack.iter().rev() {
             if var.name == name && var.active {
@@ -190,12 +248,7 @@ pub(crate) struct Compiler {
 macro_rules! push_frame {
     ($compiler:expr, $block:expr, $code:tt) => {
         {
-            $compiler.frames.push(Frame {
-                stack: Vec::new(),
-                upvalues: Vec::new(),
-                scope: 0,
-                variables_below: $compiler.frame().variables_below + $compiler.stack().len(),
-            });
+            $compiler.frames.push(Frame::new());
 
             // Return value stored as a variable
             $compiler.define_variable("", Type::Unknown, &mut $block).unwrap();
@@ -241,12 +294,7 @@ impl Compiler {
             tokens,
             current_file: PathBuf::from(current_file),
 
-            frames: vec![Frame {
-                stack: Vec::new(),
-                upvalues: Vec::new(),
-                scope: 0,
-                variables_below: 0,
-            }],
+            frames: vec![Frame::new()],
 
             panic: false,
             errors: vec![],
@@ -946,6 +994,7 @@ impl Compiler {
         expect!(self, Token::For, "Expected 'for' at start of for-loop.");
 
         push_scope!(self, block, {
+            self.frame_mut().push_loop();
             // Definition
             match self.peek_four() {
                 // TODO(ed): Typed definitions aswell!
@@ -981,6 +1030,8 @@ impl Compiler {
 
             block.patch(Op::JmpFalse(block.curr()), cond_out);
 
+            let stacksize = self.frame().stack.len();
+            self.frame_mut().pop_loop(block, stacksize, inc, block.curr());
         });
     }
 
@@ -1223,6 +1274,24 @@ impl Compiler {
                 self.for_loop(block);
             }
 
+            (Token::Break, ..) => {
+                self.eat();
+                let addr = add_op(self, block, Op::Illegal);
+                let stack_size = self.frame().stack.len();
+                if self.frame_mut().add_break(addr, stack_size, block).is_err() {;
+                    error!(self, "Cannot place 'break' outside of loop.");
+                }
+            }
+
+            (Token::Continue, ..) => {
+                self.eat();
+                let addr = add_op(self, block, Op::Illegal);
+                let stack_size = self.frame().stack.len();
+                if self.frame_mut().add_continue(addr, stack_size, block).is_err() {
+                    error!(self, "Cannot place 'continue' outside of loop.");
+                }
+            }
+
             (Token::Ret, ..) => {
                 self.eat();
                 self.expression(block);
diff --git a/src/lib.rs b/src/lib.rs
index 22b6713..ee176e5 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -397,6 +397,14 @@ pub enum Op {
     ///
     /// {A} - JmpFalse(n) - {}
     JmpFalse(usize),
+    /// Sets the instruction pointer
+    /// to the given value and pops
+    /// the given amount of values.
+    ///
+    /// Used for 'break' and 'continue'.
+    ///
+    /// {A, B, C} - JmpNPop(n, 2) - {A}
+    JmpNPop(usize, usize),
 
     /// Compares the two topmost elements
     /// on the stack for equality, and pushes
@@ -1070,6 +1078,63 @@ a.a <=> 0"
     );
 
     test_multiple!(
+        break_and_continue,
+        simple_break: "
+a := 0
+for i := 0, i < 10, i += 1 {
+    a = a + 1
+    if i == 2 {
+        break
+    }
+}
+a <=> 3
+",
+
+        simple_continue: "
+a := 0
+for i := 0, i < 4, i += 1 {
+    if i == 2 {
+        continue
+    }
+    a = a + 1
+}
+a <=> 3
+",
+
+        advanced_break: "
+a := 0
+for i := 0, i < 10, i += 1 {
+    q := 0
+    qq := 0
+    qqq := 0
+    qqqq := 0
+
+    a = a + 1
+    if i == 2 {
+        break
+    }
+}
+a <=> 3
+",
+
+        advanced_continue: "
+a := 0
+for i := 0, i < 4, i += 1 {
+    q := 0
+    qq := 0
+    qqq := 0
+    qqqq := 0
+
+    if i == 2 {
+        continue
+    }
+    a = a + 1
+}
+a <=> 3
+",
+    );
+
+    test_multiple!(
         read_constants,
         simple: "
 a :: 1
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index cf7a14c..b54e194 100644
--- a/src/tokenizer.rs
+++ b/src/tokenizer.rs
@@ -24,6 +24,10 @@ pub enum Token {
     Else,
     #[token("for")]
     For,
+    #[token("break")]
+    Break,
+    #[token("continue")]
+    Continue,
     // #[token("in")]
     // In,
     // #[token("loop")]
diff --git a/src/vm.rs b/src/vm.rs
index d238087..c4d72a1 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -336,6 +336,20 @@ impl VM {
                 }
             }
 
+            Op::JmpNPop(line, to_pop) => {
+                let hi = self.stack.len();
+                let lo = hi - to_pop;
+                for slot in lo..hi {
+                    if self.upvalues.contains_key(&slot) {
+                        let value = self.stack[slot].clone();
+                        self.drop_upvalue(slot, value);
+                    }
+                }
+                self.stack.truncate(lo);
+                self.frame_mut().ip = line;
+                return Ok(OpResult::Continue);
+            }
+
             Op::Assert => {
                 if matches!(self.pop(), Value::Bool(false)) {
                     error!(self, ErrorKind::AssertFailed);
@@ -697,6 +711,9 @@ impl VM {
                     a => { error!(self, ErrorKind::TypeError(op, vec![a.into()])) },
                 }
             }
+
+            Op::JmpNPop(_, _) => {}
+
             _ => {
                 self.eval_op(op)?;
                 return Ok(())