From da89f189999402ed4117ca4465037d5684fb6086 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 19 Feb 2020 11:24:31 +0300
Subject: [PATCH] compiler: support break and continue in range loops

---
 pkg/compiler/codegen.go  | 15 +++++++++--
 pkg/compiler/for_test.go | 58 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 04b04933c..cdb54a938 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -763,8 +763,14 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			return nil
 		}
 
-		start := c.newLabel()
-		end := c.newLabel()
+		start, label := c.generateLabel(labelStart)
+		end := c.newNamedLabel(labelEnd, label)
+		post := c.newNamedLabel(labelPost, label)
+
+		lastFor := c.currentFor
+		lastSwitch := c.currentSwitch
+		c.currentFor = label
+		c.currentSwitch = label
 
 		ast.Walk(c, n.X)
 
@@ -787,11 +793,16 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 
 		ast.Walk(c, n.Body)
 
+		c.setLabel(post)
+
 		emit.Opcode(c.prog.BinWriter, opcode.INC)
 		emit.Jmp(c.prog.BinWriter, opcode.JMP, int16(start))
 
 		c.setLabel(end)
 
+		c.currentFor = lastFor
+		c.currentSwitch = lastSwitch
+
 		return nil
 
 	// We dont really care about assertions for the core logic.
diff --git a/pkg/compiler/for_test.go b/pkg/compiler/for_test.go
index c21051951..cb9f1f536 100644
--- a/pkg/compiler/for_test.go
+++ b/pkg/compiler/for_test.go
@@ -616,6 +616,64 @@ func TestForLoopNestedContinueLabel(t *testing.T) {
 	eval(t, src, big.NewInt(15))
 }
 
+func TestForLoopRangeBreak(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		var i int
+		arr := []int{1, 2, 3}
+		for i = range arr {
+			if arr[i] == 2 {
+				break
+			}
+		}
+		return i
+	}`
+
+	eval(t, src, big.NewInt(1))
+}
+
+func TestForLoopRangeNestedBreak(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		k := 5
+		arr := []int{1, 2, 3}
+		urr := []int{4, 5, 6, 7}
+		loop:
+		for range arr {
+			k++
+			for j := range urr {
+				k++
+				if j == 3 {
+					break loop
+				}
+			}
+		}
+		return k
+	}`
+
+	eval(t, src, big.NewInt(10))
+}
+
+func TestForLoopRangeContinue(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		i := 6
+		arr := []int{1, 2, 3}
+		for j := range arr {
+			if arr[j] < 2 {
+				continue
+			}
+			i++
+		}
+		return i
+	}`
+
+	eval(t, src, big.NewInt(8))
+}
+
 func TestForLoopRangeNoVariable(t *testing.T) {
 	src := `
 	package foo