internal/lsp: move package selection to before type checking

This change moves package selection to before type checking so we don't
unnecessarily type-check both variants of a package. As a result, exec
time and memory usage for features making calls to GetParsedFile are cut
by half since we only type check either the narrowest or the widest
package.

Change-Id: Ifd076f8c38e33de2bd3509fe17feafccd59d7419
Reviewed-on: https://go-review.googlesource.com/c/tools/+/257240
Run-TryBot: Danish Dua <[email protected]>
gopls-CI: kokoro <[email protected]>
TryBot-Result: Go Bot <[email protected]>
Trust: Danish Dua <[email protected]>
Reviewed-by: Heschi Kreinick <[email protected]>
(cherry picked from commit 8d73f17)
Reviewed-on: https://go-review.googlesource.com/c/tools/+/257601
Trust: Rebecca Stambler <[email protected]>
Run-TryBot: Rebecca Stambler <[email protected]>
Reviewed-by: Danish Dua <[email protected]>

Author: dandua98

internal/lsp/cache/snapshot.go

@@ -287,6 +287,54 @@ func hashUnsavedOverlays(files map[span.URI]source.VersionedFileHandle) string {
 func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI, mode source.TypecheckMode) ([]source.Package, error) {
 	ctx = event.Label(ctx, tag.URI.Of(uri))
 
+	phs, err := s.packageHandlesForFile(ctx, uri, mode)
+	if err != nil {
+		return nil, err
+	}
+	var pkgs []source.Package
+	for _, ph := range phs {
+		pkg, err := ph.check(ctx, s)
+		if err != nil {
+			return nil, err
+		}
+		pkgs = append(pkgs, pkg)
+	}
+	return pkgs, nil
+}
+
+func (s *snapshot) PackageForFile(ctx context.Context, uri span.URI, mode source.TypecheckMode, pkgPolicy source.PackageFilter) (source.Package, error) {
+	ctx = event.Label(ctx, tag.URI.Of(uri))
+
+	phs, err := s.packageHandlesForFile(ctx, uri, mode)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(phs) < 1 {
+		return nil, errors.Errorf("no packages")
+	}
+
+	ph := phs[0]
+	for _, handle := range phs[1:] {
+		switch pkgPolicy {
+		case source.WidestPackage:
+			if ph == nil || len(handle.CompiledGoFiles()) > len(ph.CompiledGoFiles()) {
+				ph = handle
+			}
+		case source.NarrowestPackage:
+			if ph == nil || len(handle.CompiledGoFiles()) < len(ph.CompiledGoFiles()) {
+				ph = handle
+			}
+		}
+	}
+	if ph == nil {
+		return nil, errors.Errorf("no packages in input")
+	}
+
+	return ph.check(ctx, s)
+}
+
+func (s *snapshot) packageHandlesForFile(ctx context.Context, uri span.URI, mode source.TypecheckMode) ([]*packageHandle, error) {
 	// Check if we should reload metadata for the file. We don't invalidate IDs
 	// (though we should), so the IDs will be a better source of truth than the
 	// metadata. If there are no IDs for the file, then we should also reload.
@@ -310,7 +358,7 @@ func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI, mode sourc
 		}
 	}
 	// Get the list of IDs from the snapshot again, in case it has changed.
-	var pkgs []source.Package
+	var phs []*packageHandle
 	for _, id := range s.getIDsForURI(uri) {
 		var parseModes []source.ParseMode
 		switch mode {
@@ -327,22 +375,15 @@ func (s *snapshot) PackagesForFile(ctx context.Context, uri span.URI, mode sourc
 		}
 
 		for _, parseMode := range parseModes {
-			pkg, err := s.checkedPackage(ctx, id, parseMode)
+			ph, err := s.buildPackageHandle(ctx, id, parseMode)
 			if err != nil {
 				return nil, err
 			}
-			pkgs = append(pkgs, pkg)
+			phs = append(phs, ph)
 		}
 	}
-	return pkgs, nil
-}
 
-func (s *snapshot) checkedPackage(ctx context.Context, id packageID, mode source.ParseMode) (*pkg, error) {
-	ph, err := s.buildPackageHandle(ctx, id, mode)
-	if err != nil {
-		return nil, err
-	}
-	return ph.check(ctx, s)
+	return phs, nil
 }
 
 func (s *snapshot) GetReverseDependencies(ctx context.Context, id string) ([]source.Package, error) {
@@ -366,6 +407,14 @@ func (s *snapshot) GetReverseDependencies(ctx context.Context, id string) ([]sou
 	return pkgs, nil
 }
 
+func (s *snapshot) checkedPackage(ctx context.Context, id packageID, mode source.ParseMode) (*pkg, error) {
+	ph, err := s.buildPackageHandle(ctx, id, mode)
+	if err != nil {
+		return nil, err
+	}
+	return ph.check(ctx, s)
+}
+
 // transitiveReverseDependencies populates the uris map with file URIs
 // belonging to the provided package and its transitive reverse dependencies.
 func (s *snapshot) transitiveReverseDependencies(id packageID, ids map[packageID]struct{}) {

internal/lsp/code_action.go

@@ -123,11 +123,7 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
-		pkgs, err := snapshot.PackagesForFile(ctx, fh.URI(), source.TypecheckFull)
-		if err != nil {
-			return nil, err
-		}
-		pkg, err := source.WidestPackage(pkgs)
+		pkg, err := snapshot.PackageForFile(ctx, fh.URI(), source.TypecheckFull, source.WidestPackage)
 		if err != nil {
 			return nil, err
 		}

internal/lsp/link.go

@@ -100,11 +100,7 @@ func modLinks(ctx context.Context, snapshot source.Snapshot, fh source.FileHandl
 func goLinks(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.DocumentLink, error) {
 	view := snapshot.View()
 	// We don't actually need type information, so any typecheck mode is fine.
-	pkgs, err := snapshot.PackagesForFile(ctx, fh.URI(), source.TypecheckWorkspace)
-	if err != nil {
-		return nil, err
-	}
-	pkg, err := source.WidestPackage(pkgs)
+	pkg, err := snapshot.PackageForFile(ctx, fh.URI(), source.TypecheckWorkspace, source.WidestPackage)
 	if err != nil {
 		return nil, err
 	}

internal/lsp/source/completion/completion.go

@@ -404,7 +404,7 @@ func Completion(ctx context.Context, snapshot source.Snapshot, fh source.FileHan
 		// present but no package name exists.
 		items, surrounding, innerErr := packageClauseCompletions(ctx, snapshot, fh, protoPos)
 		if innerErr != nil {
-			// return the error for getParsedFile since it's more relevant in this situation.
+			// return the error for GetParsedFile since it's more relevant in this situation.
 			return nil, nil, errors.Errorf("getting file for Completion: %w (package completions: %v)", err, innerErr)
 
 		}

internal/lsp/source/util.go

@@ -75,64 +75,17 @@ func (s MappedRange) URI() span.URI {
 }
 
 // GetParsedFile is a convenience function that extracts the Package and
-// ParsedGoFile for a File in a Snapshot. selectPackage is typically
-// Narrowest/WidestPackageHandle below.
-func GetParsedFile(ctx context.Context, snapshot Snapshot, fh FileHandle, selectPackage PackagePolicy) (Package, *ParsedGoFile, error) {
-	phs, err := snapshot.PackagesForFile(ctx, fh.URI(), TypecheckWorkspace)
-	if err != nil {
-		return nil, nil, err
-	}
-	pkg, err := selectPackage(phs)
+// ParsedGoFile for a file in a Snapshot. pkgPolicy is one of NarrowestPackage/
+// WidestPackage.
+func GetParsedFile(ctx context.Context, snapshot Snapshot, fh FileHandle, pkgPolicy PackageFilter) (Package, *ParsedGoFile, error) {
+	pkg, err := snapshot.PackageForFile(ctx, fh.URI(), TypecheckWorkspace, pkgPolicy)
 	if err != nil {
 		return nil, nil, err
 	}
 	pgh, err := pkg.File(fh.URI())
 	return pkg, pgh, err
 }
 
-type PackagePolicy func([]Package) (Package, error)
-
-// NarrowestPackage picks the "narrowest" package for a given file.
-//
-// By "narrowest" package, we mean the package with the fewest number of files
-// that includes the given file. This solves the problem of test variants,
-// as the test will have more files than the non-test package.
-func NarrowestPackage(pkgs []Package) (Package, error) {
-	if len(pkgs) < 1 {
-		return nil, errors.Errorf("no packages")
-	}
-	result := pkgs[0]
-	for _, handle := range pkgs[1:] {
-		if result == nil || len(handle.CompiledGoFiles()) < len(result.CompiledGoFiles()) {
-			result = handle
-		}
-	}
-	if result == nil {
-		return nil, errors.Errorf("no packages in input")
-	}
-	return result, nil
-}
-
-// WidestPackage returns the Package containing the most files.
-//
-// This is useful for something like diagnostics, where we'd prefer to offer diagnostics
-// for as many files as possible.
-func WidestPackage(pkgs []Package) (Package, error) {
-	if len(pkgs) < 1 {
-		return nil, errors.Errorf("no packages")
-	}
-	result := pkgs[0]
-	for _, handle := range pkgs[1:] {
-		if result == nil || len(handle.CompiledGoFiles()) > len(result.CompiledGoFiles()) {
-			result = handle
-		}
-	}
-	if result == nil {
-		return nil, errors.Errorf("no packages in input")
-	}
-	return result, nil
-}
-
 func IsGenerated(ctx context.Context, snapshot Snapshot, uri span.URI) bool {
 	fh, err := snapshot.GetFile(ctx, uri)
 	if err != nil {

internal/lsp/source/view.go

@@ -106,6 +106,10 @@ type Snapshot interface {
 	// in mode.
 	PackagesForFile(ctx context.Context, uri span.URI, mode TypecheckMode) ([]Package, error)
 
+	// PackageForFile returns a single package that this file belongs to,
+	// checked in mode and filtered by the package policy.
+	PackageForFile(ctx context.Context, uri span.URI, mode TypecheckMode, selectPackage PackageFilter) (Package, error)
+
 	// GetActiveReverseDeps returns the active files belonging to the reverse
 	// dependencies of this file's package, checked in TypecheckWorkspace mode.
 	GetReverseDependencies(ctx context.Context, id string) ([]Package, error)
@@ -127,6 +131,23 @@ type Snapshot interface {
 	WorkspaceDirectories(ctx context.Context) []span.URI
 }
 
+// PackageFilter sets how a package is filtered out from a set of packages
+// containing a given file.
+type PackageFilter int
+
+const (
+	// NarrowestPackage picks the "narrowest" package for a given file.
+	// By "narrowest" package, we mean the package with the fewest number of
+	// files that includes the given file. This solves the problem of test
+	// variants, as the test will have more files than the non-test package.
+	NarrowestPackage PackageFilter = iota
+
+	// WidestPackage returns the Package containing the most files.
+	// This is useful for something like diagnostics, where we'd prefer to
+	// offer diagnostics for as many files as possible.
+	WidestPackage
+)
+
 // View represents a single workspace.
 // This is the level at which we maintain configuration like working directory
 // and build tags.