From 03dc6f7cbbfe9dc0bce2193dcfdccdc96b4928af Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 17 Feb 2020 17:19:54 +0300
Subject: [PATCH] compiler: support continue statement in for loops

---
 pkg/compiler/codegen.go  |  6 +++-
 pkg/compiler/for_test.go | 78 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index d73cd0bcf..e7f627c52 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -55,6 +55,7 @@ type labelOffsetType byte
 const (
 	labelStart labelOffsetType = iota // labelStart is a default label type
 	labelEnd                          // labelEnd is a type for labels that are targets for break
+	labelPost                         // labelPost is a type for labels that are targets for continue
 )
 
 type labelWithType struct {
@@ -692,7 +693,8 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			end := c.getLabelOffset(labelEnd, label)
 			emit.Jmp(c.prog.BinWriter, opcode.JMP, int16(end))
 		case token.CONTINUE:
-			c.prog.Err = fmt.Errorf("continue statement is not supported yet")
+			post := c.getLabelOffset(labelPost, label)
+			emit.Jmp(c.prog.BinWriter, opcode.JMP, int16(post))
 		}
 
 		return nil
@@ -707,6 +709,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 	case *ast.ForStmt:
 		fstart, label := c.generateLabel(labelStart)
 		fend := c.newNamedLabel(labelEnd, label)
+		fpost := c.newNamedLabel(labelPost, label)
 
 		lastLabel := c.currentFor
 		c.currentFor = label
@@ -725,6 +728,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 
 		// Walk body followed by the iterator (post stmt).
 		ast.Walk(c, n.Body)
+		c.setLabel(fpost)
 		if n.Post != nil {
 			ast.Walk(c, n.Post)
 		}
diff --git a/pkg/compiler/for_test.go b/pkg/compiler/for_test.go
index b0da885d9..c21051951 100644
--- a/pkg/compiler/for_test.go
+++ b/pkg/compiler/for_test.go
@@ -538,6 +538,84 @@ func TestForLoopNestedBreakLabel(t *testing.T) {
 	eval(t, src, big.NewInt(5))
 }
 
+func TestForLoopContinue(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		var i, j int
+		for i < 10 {
+			i++
+			if i >= 5 {
+				continue
+			}
+			j++
+		}
+		return j
+	}`
+
+	eval(t, src, big.NewInt(4))
+}
+
+func TestForLoopContinueLabel(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		var i, j int
+		loop:
+		for i < 10 {
+			i++
+			if i >= 5 {
+				continue loop
+			}
+			j++
+		}
+		return j
+	}`
+
+	eval(t, src, big.NewInt(4))
+}
+
+func TestForLoopNestedContinue(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		var i, k int
+		for i < 10 {
+			i++
+			for j := 0; j < 3; j++ {
+				if j >= 2 {
+					continue
+				}
+				k++
+			}
+		}
+		return k
+	}`
+
+	eval(t, src, big.NewInt(20))
+}
+
+func TestForLoopNestedContinueLabel(t *testing.T) {
+	src := `
+	package foo
+	func Main() int {
+		var i int
+		loop:
+		for ; i < 10; i += 10 {
+			i++
+			for j := 0; j < 4; j++ {
+				if i == 5 {
+					continue loop
+				}
+				i++
+			}
+		}
+		return i
+	}`
+
+	eval(t, src, big.NewInt(15))
+}
+
 func TestForLoopRangeNoVariable(t *testing.T) {
 	src := `
 	package foo