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

compiler: allow to call methods on return values

pkg/compiler/codegen.go

@@ -918,7 +918,7 @@ func (c *codegen) Visit(node ast.Node) ast.Visitor {
 
 			f, ok = c.funcs[name]
 			if ok {
-				f.selector = fun.X.(*ast.Ident)
+				f.selector = fun.X
 				isBuiltin = isCustomBuiltin(f)
 				if canInline(f.pkg.Path(), f.decl.Name.Name) {
 					c.inlineCall(f, n)
@@ -2069,15 +2069,16 @@ func (c *codegen) getFuncFromIdent(fun *ast.Ident) (*funcScope, bool) {
 // getFuncNameFromSelector returns fully-qualified function name from the selector expression.
 // Second return value is true iff this was a method call, not foreign package call.
 func (c *codegen) getFuncNameFromSelector(e *ast.SelectorExpr) (string, bool) {
-	ident := e.X.(*ast.Ident)
 	if c.typeInfo.Selections[e] != nil {
-		typ := c.typeInfo.Types[ident].Type.String()
+		typ := c.typeInfo.Types[e.X].Type.String()
 		name := c.getIdentName(typ, e.Sel.Name)
 		if name[0] == '*' {
 			name = name[1:]
 		}
 		return name, true
 	}
+
+	ident := e.X.(*ast.Ident)
 	return c.getIdentName(ident.Name, e.Sel.Name), false
 }
 

pkg/compiler/func_scope.go

@@ -13,7 +13,7 @@ type funcScope struct {
 
 	// Selector of the function if there is any. Only functions imported
 	// from other packages should have a selector.
-	selector *ast.Ident
+	selector ast.Expr
 
 	// The declaration of the function in the AST. Nil if this scope is not a function.
 	decl *ast.FuncDecl

pkg/compiler/function_call_test.go

@@ -11,6 +11,27 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestReturnValueReceiver(t *testing.T) {
+	t.Run("regular", func(t *testing.T) {
+		src := `package foo
+		import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/method"
+
+		func Main() int {
+			return method.NewX().GetA()
+		}`
+		eval(t, src, big.NewInt(42))
+	})
+	t.Run("inline", func(t *testing.T) {
+		src := `package foo
+		import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/inline"
+
+		func Main() int {
+			return inline.NewT().GetN()
+		}`
+		eval(t, src, big.NewInt(42))
+	})
+}
+
 func TestSimpleFunctionCall(t *testing.T) {
 	src := `
 		package testcase

pkg/compiler/testdata/inline/inline.go

@@ -56,3 +56,11 @@ func (t *T) Inc(i int) int {
 	t.N += i
 	return n
 }
+
+func NewT() T {
+	return T{N: 42}
+}
+
+func (t T) GetN() int {
+	return t.N
+}

pkg/compiler/testdata/method/struct.go

@@ -0,0 +1,16 @@
+package method
+
+// X is some type.
+type X struct {
+	a int
+}
+
+// GetA returns the value of a.
+func (x X) GetA() int {
+	return x.a
+}
+
+// NewX creates a new X instance.
+func NewX() X {
+	return X{a: 42}
+}