From ccb53414f25a4240184e2308973b23109eaa3fd7 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 18 Feb 2020 17:07:38 +0300
Subject: [PATCH] compiler: support break in switch statements

---
 pkg/compiler/codegen.go     | 18 +++++++++--
 pkg/compiler/switch_test.go | 60 +++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index e7f627c52..04b04933c 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -43,6 +43,8 @@ type codegen struct {
 
 	// A label for the for-loop being currently visited.
 	currentFor string
+	// A label for the switch statement being visited.
+	currentSwitch string
 	// A label to be used in the next statement.
 	nextLabel string
 
@@ -403,7 +405,10 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		ast.Walk(c, n.Tag)
 
 		eqOpcode := c.getEqualityOpcode(n.Tag)
-		switchEnd := c.newLabel()
+		switchEnd, label := c.generateLabel(labelEnd)
+
+		lastSwitch := c.currentSwitch
+		c.currentSwitch = label
 
 		for i := range n.Body.List {
 			lEnd := c.newLabel()
@@ -434,6 +439,8 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		c.setLabel(switchEnd)
 		emit.Opcode(c.prog.BinWriter, opcode.DROP)
 
+		c.currentSwitch = lastSwitch
+
 		return nil
 
 	case *ast.BasicLit:
@@ -683,9 +690,13 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.BranchStmt:
-		label := c.currentFor
+		var label string
 		if n.Label != nil {
 			label = n.Label.Name
+		} else if n.Tok == token.BREAK {
+			label = c.currentSwitch
+		} else if n.Tok == token.CONTINUE {
+			label = c.currentFor
 		}
 
 		switch n.Tok {
@@ -712,7 +723,9 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		fpost := c.newNamedLabel(labelPost, label)
 
 		lastLabel := c.currentFor
+		lastSwitch := c.currentSwitch
 		c.currentFor = label
+		c.currentSwitch = label
 
 		// Walk the initializer and condition.
 		if n.Init != nil {
@@ -738,6 +751,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		c.setLabel(fend)
 
 		c.currentFor = lastLabel
+		c.currentSwitch = lastSwitch
 
 		return nil
 
diff --git a/pkg/compiler/switch_test.go b/pkg/compiler/switch_test.go
index 8b39c3c8b..8ffe00b8a 100644
--- a/pkg/compiler/switch_test.go
+++ b/pkg/compiler/switch_test.go
@@ -127,6 +127,66 @@ var switchTestCases = []testCase{
 		}`,
 		big.NewInt(4),
 	},
+	{
+		"break from switch",
+		`package main
+		func Main() int {
+			i := 3
+			switch i {
+			case 2: return 2
+			case 3:
+				i = 1
+				break
+				return 3
+			case 4: return 4
+			}
+			return i
+		}`,
+		big.NewInt(1),
+	},
+	{
+		"break from outer for",
+		`package main
+		func Main() int {
+			i := 3
+			loop:
+			for i < 10 {
+				i++
+				switch i {
+				case 5:
+					i = 7
+					break loop
+					return 3
+				case 6: return 4
+				}
+			}
+			return i
+		}`,
+		big.NewInt(7),
+	},
+	{
+		"continue outer for",
+		`package main
+		func Main() int {
+			i := 2
+			for i < 10 {
+				i++
+				switch i {
+				case 3:
+					i = 7
+					continue
+				case 4, 5, 6, 7: return 5
+				case 8: return 2
+				}
+
+				if i == 7 {
+					return 6
+				}
+			}
+			return i
+		}`,
+		big.NewInt(2),
+	},
 }
 
 func TestSwitch(t *testing.T) {