Merge pull request #1785 from nspcc-dev/fix/compiler

Process `return` in inlined functions properly.

pkg/compiler/codegen.go

@@ -49,6 +49,9 @@ type codegen struct {
 	labels map[labelWithType]uint16
 	// A list of nested label names together with evaluation stack depth.
 	labelList []labelWithStackSize
+	// inlineLabelOffsets contains size of labelList at the start of inline call processing.
+	// For such calls we need to drop only newly created part of stack.
+	inlineLabelOffsets []int
 
 	// A label for the for-loop being currently visited.
 	currentFor string
@@ -607,7 +610,11 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 		c.setLabel(l)
 
 		cnt := 0
-		for i := range c.labelList {
+		start := 0
+		if len(c.inlineLabelOffsets) > 0 {
+			start = c.inlineLabelOffsets[len(c.inlineLabelOffsets)-1]
+		}
+		for i := start; i < len(c.labelList); i++ {
 			cnt += c.labelList[i].sz
 		}
 		c.dropItems(cnt)

pkg/compiler/inline.go

@@ -16,6 +16,14 @@ import (
 //    <inline body of f directly>
 // }
 func (c *codegen) inlineCall(f *funcScope, n *ast.CallExpr) {
+	labelSz := len(c.labelList)
+	offSz := len(c.inlineLabelOffsets)
+	c.inlineLabelOffsets = append(c.inlineLabelOffsets, labelSz)
+	defer func() {
+		c.inlineLabelOffsets = c.inlineLabelOffsets[:offSz]
+		c.labelList = c.labelList[:labelSz]
+	}()
+
 	pkg := c.buildInfo.program.Package(f.pkg.Path())
 	sig := c.typeOf(n.Fun).(*types.Signature)
 
@@ -31,7 +39,8 @@ func (c *codegen) inlineCall(f *funcScope, n *ast.CallExpr) {
 	// while stored in the new.
 	oldScope := c.scope.vars.locals
 	c.scope.vars.newScope()
-	newScope := c.scope.vars.locals
+	newScope := make([]map[string]varInfo, len(c.scope.vars.locals))
+	copy(newScope, c.scope.vars.locals)
 	defer c.scope.vars.dropScope()
 
 	hasVarArgs := !n.Ellipsis.IsValid()

pkg/compiler/inline_test.go

@@ -115,6 +115,71 @@ func TestInline(t *testing.T) {
 	})
 }
 
+func TestInlineInLoop(t *testing.T) {
+	t.Run("simple", func(t *testing.T) {
+		src := `package foo
+		import "github.com/nspcc-dev/neo-go/pkg/interop/binary"
+		import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/inline"
+		func Main() int {
+			sum := 0
+			values := []int{10, 11}
+			for _, v := range values {
+				binary.Itoa(v, 10)
+				sum += inline.VarSum(1, 2, 3, 4)
+			}
+			return sum
+		}`
+		eval(t, src, big.NewInt(20))
+	})
+	t.Run("inlined argument", func(t *testing.T) {
+		src := `package foo
+		import "github.com/nspcc-dev/neo-go/pkg/interop/binary"
+		import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/inline"
+		func Main() int {
+			sum := 0
+			values := []int{10, 11}
+			for _, v := range values {
+				binary.Itoa(v, 10)
+				sum += inline.VarSum(1, 2, 3, binary.Atoi("4", 10))
+			}
+			return sum
+		}`
+		eval(t, src, big.NewInt(20))
+	})
+	t.Run("check clean stack on return", func(t *testing.T) {
+		src := `package foo
+		import "github.com/nspcc-dev/neo-go/pkg/interop/binary"
+		import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/inline"
+		func Main() int {
+			values := []int{10, 11, 12}
+			for _, v := range values {
+				binary.Itoa(v, 10)
+				if v == 11 {
+					return inline.VarSum(2, 20, 200)
+				}
+			}
+			return 0
+		}`
+		eval(t, src, big.NewInt(222))
+	})
+}
+
+func TestInlineInSwitch(t *testing.T) {
+	src := `package foo
+	import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/inline"
+	func Main() int {
+		switch inline.VarSum(1, 2) {
+		case inline.VarSum(3, 1):
+			return 10
+		case inline.VarSum(4, -1):
+			return 11
+		default:
+			return 12
+		}
+	}`
+	eval(t, src, big.NewInt(11))
+}
+
 func TestInlineGlobalVariable(t *testing.T) {
 	t.Run("simple", func(t *testing.T) {
 		src := `package foo

pkg/compiler/vm_test.go

@@ -3,6 +3,8 @@ package compiler_test
 import (
 	"errors"
 	"fmt"
+	"math/big"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -107,6 +109,8 @@ func newStoragePlugin() *storagePlugin {
 	s.interops[interopnames.ToID([]byte(interopnames.SystemStoragePut))] = s.Put
 	s.interops[interopnames.ToID([]byte(interopnames.SystemStorageGetContext))] = s.GetContext
 	s.interops[interopnames.ToID([]byte(interopnames.SystemRuntimeNotify))] = s.Notify
+	s.interops[interopnames.ToID([]byte(interopnames.SystemBinaryAtoi))] = s.Atoi
+	s.interops[interopnames.ToID([]byte(interopnames.SystemBinaryItoa))] = s.Itoa
 	return s
 
 }
@@ -122,6 +126,24 @@ func (s *storagePlugin) syscallHandler(v *vm.VM, id uint32) error {
 	return errors.New("syscall not found")
 }
 
+func (s *storagePlugin) Atoi(v *vm.VM) error {
+	str := v.Estack().Pop().String()
+	base := v.Estack().Pop().BigInt().Int64()
+	n, err := strconv.ParseInt(str, int(base), 64)
+	if err != nil {
+		return err
+	}
+	v.Estack().PushVal(big.NewInt(n))
+	return nil
+}
+
+func (s *storagePlugin) Itoa(v *vm.VM) error {
+	n := v.Estack().Pop().BigInt()
+	base := v.Estack().Pop().BigInt()
+	v.Estack().PushVal(n.Text(int(base.Int64())))
+	return nil
+}
+
 func (s *storagePlugin) Notify(v *vm.VM) error {
 	name := v.Estack().Pop().String()
 	item := stackitem.NewArray(v.Estack().Pop().Array())