From 6e560c6c9f06e90486c4713f5c1627facba17711 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 8 Feb 2021 13:51:25 +0300
Subject: [PATCH] compiler: allow to inline var arg functions

---
 pkg/compiler/codegen.go                | 15 +++++++++++----
 pkg/compiler/inline.go                 | 22 ++++++++++++++++++++++
 pkg/compiler/inline_test.go            | 16 ++++++++++++++++
 pkg/compiler/testdata/inline/inline.go |  8 ++++++++
 4 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 5203e9af6..41b6aa963 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -892,10 +892,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 			typ, ok := c.typeOf(n.Fun).(*types.Signature)
 			if ok && typ.Variadic() && !n.Ellipsis.IsValid() {
 				// pack variadic args into an array only if last argument is not of form `...`
-				varSize := len(n.Args) - typ.Params().Len() + 1
-				c.emitReverse(varSize)
-				emit.Int(c.prog.BinWriter, int64(varSize))
-				emit.Opcodes(c.prog.BinWriter, opcode.PACK)
+				varSize := c.packVarArgs(n, typ)
 				numArgs -= varSize - 1
 			}
 			c.emitReverse(numArgs)
@@ -1232,6 +1229,16 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 	return c
 }
 
+// packVarArgs packs variadic arguments into an array
+// and returns amount of arguments packed.
+func (c *codegen) packVarArgs(n *ast.CallExpr, typ *types.Signature) int {
+	varSize := len(n.Args) - typ.Params().Len() + 1
+	c.emitReverse(varSize)
+	emit.Int(c.prog.BinWriter, int64(varSize))
+	emit.Opcodes(c.prog.BinWriter, opcode.PACK)
+	return varSize
+}
+
 // processDefers emits code for `defer` statements.
 // TRY-related opcodes handle exception as follows:
 // 1. CATCH block is executed only if exception has occurred.
diff --git a/pkg/compiler/inline.go b/pkg/compiler/inline.go
index 55e0c8364..6d2ba8cf1 100644
--- a/pkg/compiler/inline.go
+++ b/pkg/compiler/inline.go
@@ -25,8 +25,16 @@ func (c *codegen) inlineCall(f *funcScope, n *ast.CallExpr) {
 	c.scope.vars.newScope()
 	newScope := c.scope.vars.locals
 	defer c.scope.vars.dropScope()
+
+	hasVarArgs := !n.Ellipsis.IsValid()
+	needPack := sig.Variadic() && hasVarArgs
 	for i := range n.Args {
 		c.scope.vars.locals = oldScope
+		// true if normal arg or var arg is `slice...`
+		needStore := i < sig.Params().Len()-1 || !sig.Variadic() || !hasVarArgs
+		if !needStore {
+			break
+		}
 		name := sig.Params().At(i).Name()
 		if tv := c.typeAndValueOf(n.Args[i]); tv.Value != nil {
 			c.scope.vars.locals = newScope
@@ -53,6 +61,20 @@ func (c *codegen) inlineCall(f *funcScope, n *ast.CallExpr) {
 		c.emitStoreVar("", name)
 	}
 
+	if needPack {
+		// traverse variadic args and pack them
+		// if they are provided directly i.e. without `...`
+		c.scope.vars.locals = oldScope
+		for i := sig.Params().Len() - 1; i < len(n.Args); i++ {
+			ast.Walk(c, n.Args[i])
+		}
+		c.scope.vars.locals = newScope
+		c.packVarArgs(n, sig)
+		name := sig.Params().At(sig.Params().Len() - 1).Name()
+		c.scope.newLocal(name)
+		c.emitStoreVar("", name)
+	}
+
 	c.pkgInfoInline = append(c.pkgInfoInline, pkg)
 	oldMap := c.importMap
 	c.fillImportMap(f.file, pkg.Pkg)
diff --git a/pkg/compiler/inline_test.go b/pkg/compiler/inline_test.go
index 8870f1a47..ed9c16b3e 100644
--- a/pkg/compiler/inline_test.go
+++ b/pkg/compiler/inline_test.go
@@ -91,6 +91,22 @@ func TestInline(t *testing.T) {
 		checkCallCount(t, src, 1, 2)
 		eval(t, src, big.NewInt(51))
 	})
+	t.Run("var args, empty", func(t *testing.T) {
+		src := fmt.Sprintf(srcTmpl, `return inline.VarSum(11)`)
+		checkCallCount(t, src, 0, 1)
+		eval(t, src, big.NewInt(11))
+	})
+	t.Run("var args, direct", func(t *testing.T) {
+		src := fmt.Sprintf(srcTmpl, `return inline.VarSum(11, 14, 17)`)
+		checkCallCount(t, src, 0, 1)
+		eval(t, src, big.NewInt(42))
+	})
+	t.Run("var args, array", func(t *testing.T) {
+		src := fmt.Sprintf(srcTmpl, `arr := []int{14, 17} 
+			return inline.VarSum(11, arr...)`)
+		checkCallCount(t, src, 0, 1)
+		eval(t, src, big.NewInt(42))
+	})
 }
 
 func TestInlineConversion(t *testing.T) {
diff --git a/pkg/compiler/testdata/inline/inline.go b/pkg/compiler/testdata/inline/inline.go
index d44ec461a..13201acaf 100644
--- a/pkg/compiler/testdata/inline/inline.go
+++ b/pkg/compiler/testdata/inline/inline.go
@@ -30,3 +30,11 @@ func DropInsideInline() int {
 	sum(3, 4)
 	return 7
 }
+
+func VarSum(a int, b ...int) int {
+	sum := a
+	for i := range b {
+		sum += b[i]
+	}
+	return sum
+}