internal/lsp/cache: refactor and improve go get quick fixes

Generating go get quick fixes in a single place only made it harder to
get them right. Do it during diagnostics generation, as is the new
norm.

Also improve the user experience. When we fail to import a package
because one of its dependencies is missing, it makes more sense to run
go get on the package we tried to import, not the one that's missing:
that will download all of its missing dependencies if there happen to be
more.

Change-Id: Ib6a8140bccfafcb9f966d25639799dd4c7347c3d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/300072
Trust: Heschi Kreinick <[email protected]>
Run-TryBot: Heschi Kreinick <[email protected]>
Reviewed-by: Rebecca Stambler <[email protected]>
gopls-CI: kokoro <[email protected]>
TryBot-Result: Go Bot <[email protected]>

Author: heschi

internal/lsp/cache/check.go

@@ -13,7 +13,6 @@ import (
 	"go/types"
 	"path"
 	"path/filepath"
-	"regexp"
 	"sort"
 	"strings"
 	"sync"
@@ -22,7 +21,6 @@ import (
 	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/event"
-	"golang.org/x/tools/internal/lsp/command"
 	"golang.org/x/tools/internal/lsp/debug/tag"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
@@ -509,44 +507,10 @@ func typeCheck(ctx context.Context, snapshot *snapshot, m *metadata, mode source
 		return nil, err
 	}
 	pkg.diagnostics = append(pkg.diagnostics, depsErrors...)
-	if err := addGoGetFixes(ctx, snapshot, pkg); err != nil {
-		return nil, err
-	}
 
 	return pkg, nil
 }
 
-var importErrorRe = regexp.MustCompile(`could not import ([^\s]+)`)
-var missingModuleErrorRe = regexp.MustCompile(`cannot find module providing package ([^\s:]+)`)
-
-func addGoGetFixes(ctx context.Context, snapshot source.Snapshot, pkg *pkg) error {
-	if len(pkg.compiledGoFiles) == 0 || snapshot.GoModForFile(pkg.compiledGoFiles[0].URI) == "" {
-		// Go get only supports module mode for now.
-		return nil
-	}
-	for _, diag := range pkg.diagnostics {
-		matches := importErrorRe.FindStringSubmatch(diag.Message)
-		if len(matches) == 0 {
-			matches = missingModuleErrorRe.FindStringSubmatch(diag.Message)
-		}
-		if len(matches) == 0 {
-			continue
-		}
-		direct := !strings.Contains(diag.Message, "error while importing")
-		title := fmt.Sprintf("go get package %v", matches[1])
-		cmd, err := command.NewGoGetPackageCommand(title, command.GoGetPackageArgs{
-			URI:        protocol.URIFromSpanURI(pkg.compiledGoFiles[0].URI),
-			AddRequire: direct,
-			Pkg:        matches[1],
-		})
-		if err != nil {
-			return err
-		}
-		diag.SuggestedFixes = append(diag.SuggestedFixes, source.SuggestedFixFromCommand(cmd))
-	}
-	return nil
-}
-
 func (s *snapshot) depsErrors(ctx context.Context, pkg *pkg) ([]*source.Diagnostic, error) {
 	// Select packages that can't be found, and were imported in non-workspace packages.
 	// Workspace packages already show their own errors.
@@ -592,28 +556,34 @@ func (s *snapshot) depsErrors(ctx context.Context, pkg *pkg) ([]*source.Diagnost
 		}
 	}
 
-	// Apply a diagnostic to any import involved in the error, stopping after
+	// Apply a diagnostic to any import involved in the error, stopping once
 	// we reach the workspace.
 	var errors []*source.Diagnostic
 	for _, depErr := range relevantErrors {
 		for i := len(depErr.ImportStack) - 1; i >= 0; i-- {
 			item := depErr.ImportStack[i]
+			if _, ok := s.isWorkspacePackage(packageID(item)); ok {
+				break
+			}
+
 			for _, imp := range allImports[item] {
 				rng, err := source.NewMappedRange(s.FileSet(), imp.cgf.Mapper, imp.imp.Pos(), imp.imp.End()).Range()
 				if err != nil {
 					return nil, err
 				}
+				fixes, err := goGetQuickFixes(s, imp.cgf.URI, item)
+				if err != nil {
+					return nil, err
+				}
 				errors = append(errors, &source.Diagnostic{
-					URI:      imp.cgf.URI,
-					Range:    rng,
-					Severity: protocol.SeverityError,
-					Source:   source.TypeError,
-					Message:  fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
+					URI:            imp.cgf.URI,
+					Range:          rng,
+					Severity:       protocol.SeverityError,
+					Source:         source.TypeError,
+					Message:        fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
+					SuggestedFixes: fixes,
 				})
 			}
-			if _, ok := s.isWorkspacePackage(packageID(item)); ok {
-				break
-			}
 		}
 	}
 
@@ -651,12 +621,17 @@ func (s *snapshot) depsErrors(ctx context.Context, pkg *pkg) ([]*source.Diagnost
 			if err != nil {
 				return nil, err
 			}
+			fixes, err := goGetQuickFixes(s, pm.URI, item)
+			if err != nil {
+				return nil, err
+			}
 			errors = append(errors, &source.Diagnostic{
-				URI:      pm.URI,
-				Range:    rng,
-				Severity: protocol.SeverityError,
-				Source:   source.TypeError,
-				Message:  fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
+				URI:            pm.URI,
+				Range:          rng,
+				Severity:       protocol.SeverityError,
+				Source:         source.TypeError,
+				Message:        fmt.Sprintf("error while importing %v: %v", item, depErr.Err),
+				SuggestedFixes: fixes,
 			})
 			break
 		}

internal/lsp/cache/errors.go

@@ -101,6 +101,8 @@ func parseErrorDiagnostics(ctx context.Context, snapshot *snapshot, pkg *pkg, er
 	}}, nil
 }
 
+var importErrorRe = regexp.MustCompile(`could not import ([^\s]+)`)
+
 func typeErrorDiagnostics(ctx context.Context, snapshot *snapshot, pkg *pkg, e extendedError) ([]*source.Diagnostic, error) {
 	code, spn, err := typeErrorData(snapshot.FileSet(), pkg, e.primary)
 	if err != nil {
@@ -137,9 +139,33 @@ func typeErrorDiagnostics(ctx context.Context, snapshot *snapshot, pkg *pkg, e e
 			Message: secondary.Msg,
 		})
 	}
+
+	if match := importErrorRe.FindStringSubmatch(e.primary.Msg); match != nil {
+		diag.SuggestedFixes, err = goGetQuickFixes(snapshot, spn.URI(), match[1])
+		if err != nil {
+			return nil, err
+		}
+	}
 	return []*source.Diagnostic{diag}, nil
 }
 
+func goGetQuickFixes(snapshot *snapshot, uri span.URI, pkg string) ([]source.SuggestedFix, error) {
+	// Go get only supports module mode for now.
+	if snapshot.workspaceMode()&moduleMode == 0 {
+		return nil, nil
+	}
+	title := fmt.Sprintf("go get package %v", pkg)
+	cmd, err := command.NewGoGetPackageCommand(title, command.GoGetPackageArgs{
+		URI:        protocol.URIFromSpanURI(uri),
+		AddRequire: true,
+		Pkg:        pkg,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return []source.SuggestedFix{source.SuggestedFixFromCommand(cmd)}, nil
+}
+
 func analysisDiagnosticDiagnostics(ctx context.Context, snapshot *snapshot, pkg *pkg, a *analysis.Analyzer, e *analysis.Diagnostic) ([]*source.Diagnostic, error) {
 	var srcAnalyzer *source.Analyzer
 	// Find the analyzer that generated this diagnostic.