From e1a581be0e974f93983e9f6934e3c6a4c69c09b0 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 11 Jul 2022 19:28:15 +0300
Subject: [PATCH] compiler: allow to call methods on return values

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/compiler/codegen.go                |  7 ++++---
 pkg/compiler/func_scope.go             |  2 +-
 pkg/compiler/function_call_test.go     | 21 +++++++++++++++++++++
 pkg/compiler/testdata/inline/inline.go |  8 ++++++++
 pkg/compiler/testdata/method/struct.go | 16 ++++++++++++++++
 5 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 pkg/compiler/testdata/method/struct.go

diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 5375feb64..be22efec4 100644
--- a/pkg/compiler/codegen.go
+++ b/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
 }
 
diff --git a/pkg/compiler/func_scope.go b/pkg/compiler/func_scope.go
index cf076ee76..c6b6e19b7 100644
--- a/pkg/compiler/func_scope.go
+++ b/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
diff --git a/pkg/compiler/function_call_test.go b/pkg/compiler/function_call_test.go
index 5a265591b..1babb84e7 100644
--- a/pkg/compiler/function_call_test.go
+++ b/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
diff --git a/pkg/compiler/testdata/inline/inline.go b/pkg/compiler/testdata/inline/inline.go
index dfc7aae4e..80a7041d9 100644
--- a/pkg/compiler/testdata/inline/inline.go
+++ b/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
+}
diff --git a/pkg/compiler/testdata/method/struct.go b/pkg/compiler/testdata/method/struct.go
new file mode 100644
index 000000000..1df97b466
--- /dev/null
+++ b/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}
+}