compiler: make ForEachFile accept a package

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>

pkg/compiler/analysis.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
 	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
-	"golang.org/x/tools/go/loader"
 )
 
 var (
@@ -30,7 +29,7 @@ func (c *codegen) newGlobal(name string) {
 // and returns number of variables initialized.
 func (c *codegen) traverseGlobals() int {
 	var n int
-	c.ForEachFile(func(f *ast.File) {
+	c.ForEachFile(func(f *ast.File, _ *types.Package) {
 		n += countGlobals(f)
 	})
 	if n != 0 {
@@ -95,31 +94,29 @@ func lastStmtIsReturn(decl *ast.FuncDecl) (b bool) {
 	return false
 }
 
-func analyzeFuncUsage(mainPkg *loader.PackageInfo, pkgs map[*types.Package]*loader.PackageInfo) funcUsage {
+func (c *codegen) analyzeFuncUsage() funcUsage {
 	usage := funcUsage{}
 
-	for _, pkg := range pkgs {
-		isMain := pkg == mainPkg
-		for _, f := range pkg.Files {
-			ast.Inspect(f, func(node ast.Node) bool {
-				switch n := node.(type) {
-				case *ast.CallExpr:
-					switch t := n.Fun.(type) {
-					case *ast.Ident:
-						usage[t.Name] = true
-					case *ast.SelectorExpr:
-						usage[t.Sel.Name] = true
-					}
-				case *ast.FuncDecl:
-					// exported functions are always assumed to be used
-					if isMain && n.Name.IsExported() {
-						usage[n.Name.Name] = true
-					}
+	c.ForEachFile(func(f *ast.File, pkg *types.Package) {
+		isMain := pkg == c.mainPkg.Pkg
+		ast.Inspect(f, func(node ast.Node) bool {
+			switch n := node.(type) {
+			case *ast.CallExpr:
+				switch t := n.Fun.(type) {
+				case *ast.Ident:
+					usage[t.Name] = true
+				case *ast.SelectorExpr:
+					usage[t.Sel.Name] = true
 				}
-				return true
-			})
-		}
-	}
+			case *ast.FuncDecl:
+				// exported functions are always assumed to be used
+				if isMain && n.Name.IsExported() {
+					usage[n.Name.Name] = true
+				}
+			}
+			return true
+		})
+	})
 	return usage
 }
 

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.File) {
+func (c *codegen) convertGlobals(f *ast.File, _ *types.Package) {
 	ast.Inspect(f, func(node ast.Node) bool {
 		switch n := node.(type) {
 		case *ast.FuncDecl:
@@ -1415,16 +1415,12 @@ func (c *codegen) newLambda(u uint16, lit *ast.FuncLit) {
 }
 
 func (c *codegen) compile(info *buildInfo, pkg *loader.PackageInfo) error {
-	funUsage := analyzeFuncUsage(pkg, info.program.AllPackages)
+	c.mainPkg = pkg
+	funUsage := c.analyzeFuncUsage()
 
 	// Bring all imported functions into scope.
-	for _, pkg := range info.program.AllPackages {
-		for _, f := range pkg.Files {
-			c.resolveFuncDecls(f, pkg.Pkg)
-		}
-	}
+	c.ForEachFile(c.resolveFuncDecls)
 
-	c.mainPkg = pkg
 	n := c.traverseGlobals()
 	if n > 0 {
 		emit.Opcode(c.prog.BinWriter, opcode.RET)
@@ -1439,23 +1435,18 @@ func (c *codegen) compile(info *buildInfo, pkg *loader.PackageInfo) error {
 	sort.Slice(keys, func(i, j int) bool { return keys[i].Path() < keys[j].Path() })
 
 	// Generate the code for the program.
-	for _, k := range keys {
-		pkg := info.program.AllPackages[k]
-		c.typeInfo = &pkg.Info
-
-		for _, f := range pkg.Files {
-			for _, decl := range f.Decls {
-				switch n := decl.(type) {
-				case *ast.FuncDecl:
-					// Don't convert the function if it's not used. This will save a lot
-					// of bytecode space.
-					if funUsage.funcUsed(n.Name.Name) {
-						c.convertFuncDecl(f, n, k)
-					}
+	c.ForEachFile(func(f *ast.File, pkg *types.Package) {
+		for _, decl := range f.Decls {
+			switch n := decl.(type) {
+			case *ast.FuncDecl:
+				// Don't convert the function if it's not used. This will save a lot
+				// of bytecode space.
+				if funUsage.funcUsed(n.Name.Name) {
+					c.convertFuncDecl(f, n, pkg)
 				}
 			}
 		}
-	}
+	})
 
 	return c.prog.Err
 }

pkg/compiler/compiler.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"go/ast"
 	"go/parser"
+	"go/types"
 	"io"
 	"io/ioutil"
 	"os"
@@ -44,11 +45,11 @@ type buildInfo struct {
 }
 
 // ForEachFile executes fn on each file used in current program.
-func (c *codegen) ForEachFile(fn func(*ast.File)) {
+func (c *codegen) ForEachFile(fn func(*ast.File, *types.Package)) {
 	for _, pkg := range c.buildInfo.program.AllPackages {
 		c.typeInfo = &pkg.Info
 		for _, f := range pkg.Files {
-			fn(f)
+			fn(f, pkg.Pkg)
 		}
 	}
 }