optimize resolving package dirs in module mode

parser/parse.go

@@ -33,6 +33,8 @@ import (
 	"sort"
 	"strings"
 
+	"golang.org/x/tools/go/packages"
+
 	"k8s.io/gengo/types"
 	"k8s.io/klog/v2"
 )
@@ -48,6 +50,11 @@ type Builder struct {
 	// If true, include *_test.go
 	IncludeTestFiles bool
 
+	// Cache of import path to filesystem directory.
+	// Used to optimize loading of individual packages using ImportDir instead of Import.
+	// For context, see https://github.com/golang/go/issues/31087
+	pkgDirs map[string]string
+
 	// Map of package names to more canonical information about the package.
 	// This might hold the same value for multiple names, e.g. if someone
 	// referenced ./pkg/name or in the case of vendoring, which canonicalizes
@@ -104,6 +111,7 @@ func New() *Builder {
 	c.CgoEnabled = false
 	return &Builder{
 		context:               &c,
+		pkgDirs:               map[string]string{},
 		buildPackages:         map[importPathString]*build.Package{},
 		typeCheckedPackages:   map[importPathString]*tc.Package{},
 		fset:                  token.NewFileSet(),
@@ -230,10 +238,97 @@ func (b *Builder) addFile(pkgPath importPathString, path string, src []byte, use
 	return nil
 }
 
+// getResolvedPackageDir returns the resolved absolute path to the directory
+// containing the source for the specified package import path.
+// The returned dir is suitable to pass to build.Context#ImportDir.
+func (b *Builder) getResolvedPackageDir(pkgPath string) (string, bool) {
+	if len(pkgPath) == 0 || !strings.Contains(pkgPath, "/") {
+		return "", false
+	}
+	dir, ok := b.pkgDirs[pkgPath]
+	if ok {
+		return dir, len(dir) > 0
+	}
+	parentDir, ok := b.getResolvedPackageDir(path.Dir(pkgPath))
+	if !ok {
+		return "", false
+	}
+	return filepath.Join(parentDir, filepath.FromSlash(path.Base(pkgPath))), true
+}
+
+func (b *Builder) resolvePackageDirs(patterns []string) {
+	// make packages canonical before loading
+	canonicalPatterns := make([]string, 0, len(patterns))
+	for _, p := range patterns {
+		canonicalPattern := string(canonicalizeImportPath(p))
+		if _, ok := b.getResolvedPackageDir(canonicalPattern); ok {
+			// we already had a module dir resolved for this package
+			continue
+		}
+		// append the pattern
+		canonicalPatterns = append(canonicalPatterns, canonicalPattern)
+		// if this is a recursive pattern, also append the specific root to ensure it resolves (package expansion doesn't work on symlinked paths)
+		if strings.HasSuffix(canonicalPattern, "/...") {
+			canonicalPatterns = append(canonicalPatterns, strings.TrimSuffix(canonicalPattern, "/..."))
+		}
+	}
+
+	// name
+	// module info to get root dir
+	// deps to get downstream dirs to speed up type checking
+	// files to get source location
+	cfg := &packages.Config{
+		Mode: packages.NeedName | packages.NeedFiles | packages.NeedModule | packages.NeedImports | packages.NeedDeps,
+	}
+	if len(b.context.BuildTags) > 0 {
+		cfg.BuildFlags = []string{"-tags", strings.Join(b.context.BuildTags, ",")}
+	}
+	if resolvedWorkingDir, err := os.Getwd(); err != nil {
+		klog.V(2).Infof("error resolving working dir: %v", err)
+	} else if symlinkResolvedWorkingDir, err := filepath.EvalSymlinks(resolvedWorkingDir); err != nil {
+		klog.V(2).Infof("error resolving working dir: %v", err)
+	} else {
+		cfg.Dir = symlinkResolvedWorkingDir
+	}
+	pkgs, err := packages.Load(cfg, canonicalPatterns...)
+	if err != nil {
+		klog.V(2).Infof("Error resolving packages: %v", err)
+		return
+	}
+	packages.Visit(pkgs, func(pkg *packages.Package) bool {
+		if len(pkg.Errors) > 0 {
+			klog.V(2).Info(pkg.Errors)
+		}
+		if pkg.Module != nil {
+			if len(pkg.Module.Dir) > 0 {
+				b.pkgDirs[pkg.Module.Path] = pkg.Module.Dir
+			}
+			if pkg.Module.Replace != nil && len(pkg.Module.Replace.Dir) > 0 {
+				b.pkgDirs[pkg.Module.Path] = pkg.Module.Replace.Dir
+			}
+			if pkg.Module.Error != nil {
+				klog.V(2).Infof("%#v\n", *pkg.Module.Error)
+			}
+		}
+		// use any file info we got to identify the package dir
+		if len(pkg.GoFiles) > 0 {
+			b.pkgDirs[pkg.PkgPath] = filepath.Dir(pkg.GoFiles[0])
+		} else if len(pkg.OtherFiles) > 0 {
+			b.pkgDirs[pkg.PkgPath] = filepath.Dir(pkg.OtherFiles[0])
+		} else if len(pkg.IgnoredFiles) > 0 {
+			b.pkgDirs[pkg.PkgPath] = filepath.Dir(pkg.IgnoredFiles[0])
+		}
+		return true
+	}, nil)
+}
+
 // AddPackagePatterns adds the specified patterns,
 // which may be individual import paths as used in import directives,
 // or recursive paths like `example.com/...`.
 func (b *Builder) AddPackagePatterns(patterns ...string) error {
+	// resolve patterns to speed up load time
+	b.resolvePackageDirs(patterns)
+
 	for _, d := range patterns {
 		var err error
 		if strings.HasSuffix(d, "/...") {
@@ -631,7 +726,7 @@ func (b *Builder) findTypesIn(pkgPath importPathString, u *types.Universe) error
 	return nil
 }
 
-func (b *Builder) importWithMode(dir string, mode build.ImportMode) (*build.Package, error) {
+func (b *Builder) importWithMode(pkgPath string, mode build.ImportMode) (*build.Package, error) {
 	// This is a bit of a hack.  The srcDir argument to Import() should
 	// properly be the dir of the file which depends on the package to be
 	// imported, so that vendoring can work properly and local paths can
@@ -645,9 +740,19 @@ func (b *Builder) importWithMode(dir string, mode build.ImportMode) (*build.Pack
 	}
 
 	// normalize to drop /vendor/ if present
-	dir = string(canonicalizeImportPath(dir))
+	pkgPath = string(canonicalizeImportPath(pkgPath))
+
+	if pkgDir, ok := b.getResolvedPackageDir(pkgPath); ok {
+		buildPkg, err := b.context.ImportDir(pkgDir, mode)
+		if err != nil {
+			return nil, err
+		}
+		// ensure the ImportPath is the canonical one
+		buildPkg.ImportPath = pkgPath
+		return buildPkg, nil
+	}
 
-	buildPkg, err := b.context.Import(filepath.ToSlash(dir), cwd, mode)
+	buildPkg, err := b.context.Import(pkgPath, cwd, mode)
 	if err != nil {
 		return nil, err
 	}