From 9face190abb206b87186ee9ea0c55e4df3a72d77 Mon Sep 17 00:00:00 2001
From: Joey <jmhain@protonmail.com>
Date: Tue, 19 Feb 2019 19:40:34 -0800
Subject: [PATCH] Fix continue statement

---
 tests/snippets/control_flow.py | 23 +++++++++++++++++++++++
 vm/src/frame.rs                | 13 +++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/tests/snippets/control_flow.py b/tests/snippets/control_flow.py
index e25d6eee2..a0ac46613 100644
--- a/tests/snippets/control_flow.py
+++ b/tests/snippets/control_flow.py
@@ -1,3 +1,5 @@
+# break from a nested for loop
+
 def foo():
     sum = 0
     for i in range(10):
@@ -8,3 +10,24 @@ def foo():
     return sum
 
 assert foo() == 45
+
+
+# continue statement
+
+def primes(limit):
+    """Finds all the primes from 2 up to a given number using the Sieve of Eratosthenes."""
+    sieve = [False] * (limit + 1)
+    for i in range(2, limit + 1):
+        if sieve[i]:
+            continue
+        yield i
+
+        for j in range(2 * i, limit + 1, i):
+            sieve[j] = True
+
+
+assert list(primes(1)) == []
+assert list(primes(2)) == [2]
+assert list(primes(10)) == [2, 3, 5, 7]
+assert list(primes(13)) == [2, 3, 5, 7, 11, 13]
+
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
index 4c0749999..b1374b2f3 100644
--- a/vm/src/frame.rs
+++ b/vm/src/frame.rs
@@ -542,7 +542,10 @@ impl Frame {
             bytecode::Instruction::Break => {
                 let block = self.unwind_loop(vm);
                 if let BlockType::Loop { end, .. } = block.typ {
+                    self.pop_block();
                     self.jump(end);
+                } else {
+                    unreachable!()
                 }
                 Ok(None)
             }
@@ -555,7 +558,7 @@ impl Frame {
                 if let BlockType::Loop { start, .. } = block.typ {
                     self.jump(start);
                 } else {
-                    assert!(false);
+                    unreachable!();
                 }
                 Ok(None)
             }
@@ -745,7 +748,7 @@ impl Frame {
 
     fn unwind_loop(&mut self, vm: &mut VirtualMachine) -> Block {
         loop {
-            let block = self.pop_block().expect("not in a loop");
+            let block = self.current_block().cloned().expect("not in a loop");
             match block.typ {
                 BlockType::Loop { .. } => break block,
                 BlockType::TryExcept { .. } => {
@@ -760,6 +763,8 @@ impl Frame {
                     }
                 },
             }
+
+            self.pop_block();
         }
     }
 
@@ -1080,6 +1085,10 @@ impl Frame {
         Some(block)
     }
 
+    fn current_block(&self) -> Option<&Block> {
+        self.blocks.last()
+    }
+
     pub fn push_value(&mut self, obj: PyObjectRef) {
         self.stack.push(obj);
     }