From 700941732517f9c5b28c55c8151f3c20531359d3 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 28 Jul 2020 10:59:21 +0300
Subject: [PATCH] compiler: allow to declare global variables in multiple files

Traverse and count globals across all used files.

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 pkg/compiler/analysis.go             | 10 ++++------
 pkg/compiler/codegen.go              |  4 ++--
 pkg/compiler/compiler.go             | 11 +++++++++++
 pkg/compiler/global_test.go          |  9 +++++++++
 pkg/compiler/testdata/multi/file1.go |  3 +++
 pkg/compiler/testdata/multi/file2.go |  7 +++++++
 6 files changed, 36 insertions(+), 8 deletions(-)
 create mode 100644 pkg/compiler/testdata/multi/file1.go
 create mode 100644 pkg/compiler/testdata/multi/file2.go

diff --git a/pkg/compiler/analysis.go b/pkg/compiler/analysis.go
index 0033e427b..3dc20a382 100644
--- a/pkg/compiler/analysis.go
+++ b/pkg/compiler/analysis.go
@@ -28,20 +28,18 @@ func (c *codegen) newGlobal(name string) {
 
 // traverseGlobals visits and initializes global variables.
 // and returns number of variables initialized.
-func (c *codegen) traverseGlobals(fs ...*ast.File) int {
+func (c *codegen) traverseGlobals() int {
 	var n int
-	for _, f := range fs {
+	c.ForEachFile(func(f *ast.File) {
 		n += countGlobals(f)
-	}
+	})
 	if n != 0 {
 		if n > 255 {
 			c.prog.BinWriter.Err = errors.New("too many global variables")
 			return 0
 		}
 		emit.Instruction(c.prog.BinWriter, opcode.INITSSLOT, []byte{byte(n)})
-		for _, f := range fs {
-			c.convertGlobals(f)
-		}
+		c.ForEachFile(c.convertGlobals)
 	}
 	return n
 }
diff --git a/pkg/compiler/codegen.go b/pkg/compiler/codegen.go
index 5b96f5f94..f52ec6f15 100644
--- a/pkg/compiler/codegen.go
+++ b/pkg/compiler/codegen.go
@@ -252,7 +252,7 @@ func (c *codegen) emitDefault(t types.Type) {
 // convertGlobals traverses the AST and only converts global declarations.
 // If we call this in convertFuncDecl then it will load all global variables
 // into the scope of the function.
-func (c *codegen) convertGlobals(f ast.Node) {
+func (c *codegen) convertGlobals(f *ast.File) {
 	ast.Inspect(f, func(node ast.Node) bool {
 		switch n := node.(type) {
 		case *ast.FuncDecl:
@@ -1425,7 +1425,7 @@ func (c *codegen) compile(info *buildInfo, pkg *loader.PackageInfo) error {
 	}
 
 	c.mainPkg = pkg
-	n := c.traverseGlobals(pkg.Files...)
+	n := c.traverseGlobals()
 	if n > 0 {
 		emit.Opcode(c.prog.BinWriter, opcode.RET)
 		c.initEndOffset = c.prog.Len()
diff --git a/pkg/compiler/compiler.go b/pkg/compiler/compiler.go
index 4c16af02e..35d96f3a7 100644
--- a/pkg/compiler/compiler.go
+++ b/pkg/compiler/compiler.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"go/ast"
 	"go/parser"
 	"io"
 	"io/ioutil"
@@ -42,6 +43,16 @@ type buildInfo struct {
 	program        *loader.Program
 }
 
+// ForEachFile executes fn on each file used in current program.
+func (c *codegen) ForEachFile(fn func(*ast.File)) {
+	for _, pkg := range c.buildInfo.program.AllPackages {
+		c.typeInfo = &pkg.Info
+		for _, f := range pkg.Files {
+			fn(f)
+		}
+	}
+}
+
 func getBuildInfo(src interface{}) (*buildInfo, error) {
 	conf := loader.Config{ParserMode: parser.ParseComments}
 	f, err := conf.ParseFile("", src)
diff --git a/pkg/compiler/global_test.go b/pkg/compiler/global_test.go
index 359e8ef69..e49d516d2 100644
--- a/pkg/compiler/global_test.go
+++ b/pkg/compiler/global_test.go
@@ -127,3 +127,12 @@ func TestContractWithNoMain(t *testing.T) {
 	require.Equal(t, 1, v.Estack().Len())
 	require.Equal(t, big.NewInt(42), v.PopResult())
 }
+
+func TestMultipleFiles(t *testing.T) {
+	src := `package foo
+	import "github.com/nspcc-dev/neo-go/pkg/compiler/testdata/multi"
+	func Main() int {
+		return multi.Sum()
+	}`
+	eval(t, src, big.NewInt(42))
+}
diff --git a/pkg/compiler/testdata/multi/file1.go b/pkg/compiler/testdata/multi/file1.go
new file mode 100644
index 000000000..7aa330c8e
--- /dev/null
+++ b/pkg/compiler/testdata/multi/file1.go
@@ -0,0 +1,3 @@
+package multi
+
+var SomeVar12 = 12
diff --git a/pkg/compiler/testdata/multi/file2.go b/pkg/compiler/testdata/multi/file2.go
new file mode 100644
index 000000000..2ee034599
--- /dev/null
+++ b/pkg/compiler/testdata/multi/file2.go
@@ -0,0 +1,7 @@
+package multi
+
+var SomeVar30 = 30
+
+func Sum() int {
+	return SomeVar12 + SomeVar30
+}