Merge pull request #574 from nspcc-dev/feature/multireturn

compiler: implement multiple return support, closes #562.

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 {
+		// first result should be on top of the stack
-			ast.Walk(c, n.Results[0])
+		for i := len(n.Results) - 1; i >= 0; i-- {
+			ast.Walk(c, n.Results[i])
 		}
 
 		emitOpcode(c.prog.BinWriter, opcode.FROMALTSTACK)

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))
+}