From f4571ba8cfb1d6220a19ae61b7c2e0bbfd59788b Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 24 Dec 2019 16:46:43 +0300
Subject: [PATCH] compiler: implement multiple return support

---
 pkg/compiler/codegen.go     | 17 ++++-----
 pkg/compiler/return_test.go | 71 +++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 8 deletions(-)
 create mode 100644 pkg/compiler/return_test.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 04e47430a..d3c606ef6 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -232,6 +232,8 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.AssignStmt:
+		multiRet := len(n.Rhs) != len(n.Lhs)
+
 		for i := 0; i < len(n.Lhs); i++ {
 			switch t := n.Lhs[i].(type) {
 			case *ast.Ident:
@@ -243,7 +245,10 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 					l := c.scope.loadLocal(t.Name)
 					c.emitStoreLocal(l)
 				default:
-					ast.Walk(c, n.Rhs[i])
+					if i == 0 || !multiRet {
+						ast.Walk(c, n.Rhs[i])
+					}
+
 					l := c.scope.loadLocal(t.Name)
 					c.emitStoreLocal(l)
 				}
@@ -291,16 +296,12 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		return nil
 
 	case *ast.ReturnStmt:
-		if len(n.Results) > 1 {
-			c.prog.Err = fmt.Errorf("multiple returns not supported")
-			return nil
-		}
-
 		l := c.newLabel()
 		c.setLabel(l)
 
-		if len(n.Results) > 0 {
-			ast.Walk(c, n.Results[0])
+		// first result should be on top of the stack
+		for i := len(n.Results) - 1; i >= 0; i-- {
+			ast.Walk(c, n.Results[i])
 		}
 
 		emitOpcode(c.prog.BinWriter, opcode.FROMALTSTACK)
diff --git a/pkg/compiler/return_test.go b/pkg/compiler/return_test.go
new file mode 100644
index 000000000..03a8589fa
--- /dev/null
+++ b/pkg/compiler/return_test.go
@@ -0,0 +1,71 @@
+package compiler_test
+
+import (
+	"math/big"
+	"testing"
+)
+
+func TestMultipleReturn1(t *testing.T) {
+	src := `
+		package hello
+
+		func two() (int, int) {
+			return 5, 9
+		}
+
+		func Main() int {
+			a, _ := two()
+			return a
+		}
+	`
+	eval(t, src, big.NewInt(5))
+}
+
+func TestMultipleReturn2(t *testing.T) {
+	src := `
+		package hello
+
+		func two() (int, int) {
+			return 5, 9
+		}
+
+		func Main() int {
+			_, b := two()
+			return b
+		}
+	`
+	eval(t, src, big.NewInt(9))
+}
+
+func TestMultipleReturnWithArg(t *testing.T) {
+	src := `
+		package hello
+		
+		func inc2(a int) (int, int) {
+			return a+1, a+2 
+		}
+
+		func Main() int {
+			a, b := 3, 9
+			a, b = inc2(a)
+			return a+b
+		}
+	`
+	eval(t, src, big.NewInt(9))
+}
+
+func TestSingleReturn(t *testing.T) {
+	src := `
+		package hello
+
+		func inc(k int) int {
+			return k+1
+		}
+
+		func Main() int {
+			a, b := inc(3), inc(4)
+			return a+b
+		}
+	`
+	eval(t, src, big.NewInt(9))
+}