1 files changed, 72 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);