add break and continue

4 files changed, 101 insertions, 0 deletions

diff --git a/src/compiler.rs b/src/compiler.rs
index 1f3d9d8..e0b65c1 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -128,7 +128,13 @@ struct Variable {
     captured: bool,
 }
 
+enum LoopOp {
+    Continue,
+    Break,
+}
+
 struct Frame {
+    loops: Vec<Vec<(usize, usize, LoopOp)>>,
     stack: Vec<Variable>,
     upvalues: Vec<Variable>,
     scope: usize,
@@ -137,12 +143,57 @@ struct Frame {
 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 {
@@ -860,6 +911,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!
@@ -895,6 +947,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());
         });
     }
 
@@ -1132,6 +1186,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 ae45c45..67e631a 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 also pops
+    /// the given number 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
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index 9574af1..d7b28aa 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 e5c7b2e..057fdba 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -318,6 +318,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);
@@ -678,6 +692,9 @@ impl VM {
                     a => { error!(self, ErrorKind::TypeError(op, vec![a.into()])) },
                 }
             }
+
+            Op::JmpNPop(_, _) => {}
+
             _ => {
                 self.eval_op(op)?;
                 return Ok(())