From 8a0429036bfa4221c589b62f1159bb76290bdab5 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 25 Oct 2021 10:52:15 +0300
Subject: [PATCH] compiler: set type information during traversal, fix #2231

Set all necessary context before file traversal, not only import
maps. Also, we can skip restoring import maps because all our code is
processed via `For*` iterators which set necessary context.
We can also refactor this a bit to have all context in one place,
this will be done in #2086.

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/compiler/analysis.go                 |  8 ++++++--
 pkg/compiler/function_call_test.go       | 13 +++++++++++++
 pkg/compiler/testdata/nestedcall/call.go | 13 +++++++++++++
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/pkg/compiler/analysis.go b/pkg/compiler/analysis.go
index 57f2db7ba..d39549d6f 100644
--- a/pkg/compiler/analysis.go
+++ b/pkg/compiler/analysis.go
@@ -301,7 +301,12 @@ func (c *codegen) analyzeFuncUsage() funcUsage {
 			}
 			usage[name] = true
 
-			old := c.importMap
+			pkg := c.mainPkg
+			if fd.path != "" {
+				pkg = c.buildInfo.program.Package(fd.path)
+			}
+			c.typeInfo = &pkg.Info
+			c.currPkg = pkg.Pkg
 			c.importMap = fd.importMap
 			ast.Inspect(fd.decl, func(node ast.Node) bool {
 				switch n := node.(type) {
@@ -316,7 +321,6 @@ func (c *codegen) analyzeFuncUsage() funcUsage {
 				}
 				return true
 			})
-			c.importMap = old
 		}
 		diff = nextDiff
 	}
diff --git a/pkg/compiler/function_call_test.go b/pkg/compiler/function_call_test.go
index a256eb47c..9503f59ba 100644
--- a/pkg/compiler/function_call_test.go
+++ b/pkg/compiler/function_call_test.go
@@ -350,4 +350,17 @@ func TestUnusedFunctions(t *testing.T) {
 		require.NoError(t, err)
 		eval(t, src, big.NewInt(65))
 	})
+	t.Run("method inside of an imported package", func(t *testing.T) {
+		// Check that import map is set correctly during package traversal.
+		src := `package foo
+		import inner "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/nestedcall"
+		func Main() int {
+			var t inner.Token
+			return t.Method()
+		}`
+
+		_, err := compiler.Compile("foo", strings.NewReader(src))
+		require.NoError(t, err)
+		eval(t, src, big.NewInt(2231))
+	})
 }
diff --git a/pkg/compiler/testdata/nestedcall/call.go b/pkg/compiler/testdata/nestedcall/call.go
index bd05c6f21..4d306bc7b 100644
--- a/pkg/compiler/testdata/nestedcall/call.go
+++ b/pkg/compiler/testdata/nestedcall/call.go
@@ -28,3 +28,16 @@ func y() int {
 	tmp := 10
 	return tmp
 }
+
+// Token is stateless token.
+type Token struct{}
+
+// Method is a method.
+func (t Token) Method() int {
+	return t.Inner()
+}
+
+// Inner is a function to be called in Method.
+func (t Token) Inner() int {
+	return 2231
+}