internal/lsp: add quickfixes for missing dependencies in go.mod

This change adds quick fixes for diagnostics in .go files, specifically for diagnostics that deal with imported packages that are not declared in the go.mod file. These quick fixes will automatically add the dependency in the go.mod file and format the file if there are any issues.

Updates golang/go#31999

Change-Id: Iab151ce96194fae4b1995859aec416c5473da6e3
Reviewed-on: https://go-review.googlesource.com/c/tools/+/215898
Run-TryBot: Rohan Challa <rohan@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/lsp/cmd/test/suggested_fix.go

@@ -7,6 +7,7 @@ package cmdtest
 import (
 	"testing"
 
+	"golang.org/x/tools/internal/lsp/tests"
 	"golang.org/x/tools/internal/span"
 )
 
@@ -14,7 +15,7 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span) {
 	uri := spn.URI()
 	filename := uri.Filename()
 	got, _ := r.NormalizeGoplsCmd(t, "fix", "-a", filename)
-	want := string(r.data.Golden("suggestedfix", filename, func() ([]byte, error) {
+	want := string(r.data.Golden("suggestedfix_"+tests.SpanName(spn), filename, func() ([]byte, error) {
 		return []byte(got), nil
 	}))
 	if want != got {

internal/lsp/code_action.go

@@ -101,6 +101,13 @@ func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionPara
 					}
 				}
 			}
+			actions, err := mod.SuggestedGoFixes(ctx, snapshot, fh, diagnostics)
+			if err != nil {
+				log.Error(ctx, "quick fixes failed", err, telemetry.File.Of(uri))
+			}
+			if len(actions) > 0 {
+				codeActions = append(codeActions, actions...)
+			}
 		}
 		if wanted[protocol.SourceOrganizeImports] && len(edits) > 0 {
 			codeActions = append(codeActions, protocol.CodeAction{

internal/lsp/lsp_test.go

@@ -8,7 +8,6 @@ import (
 	"context"
 	"fmt"
 	"go/token"
-	"io/ioutil"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -19,7 +18,6 @@ import (
 	"golang.org/x/tools/go/packages/packagestest"
 	"golang.org/x/tools/internal/lsp/cache"
 	"golang.org/x/tools/internal/lsp/diff"
-	"golang.org/x/tools/internal/lsp/mod"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/lsp/tests"
@@ -351,12 +349,18 @@ func (r *runner) Import(t *testing.T, spn span.Span) {
 
 func (r *runner) SuggestedFix(t *testing.T, spn span.Span) {
 	uri := spn.URI()
-	filename := uri.Filename()
 	view, err := r.server.session.ViewOf(uri)
 	if err != nil {
 		t.Fatal(err)
 	}
-
+	m, err := r.data.Mapper(uri)
+	if err != nil {
+		t.Fatal(err)
+	}
+	rng, err := m.Range(spn)
+	if err != nil {
+		t.Fatal(err)
+	}
 	// Get the diagnostics for this view if we have not done it before.
 	if r.diagnostics == nil {
 		r.diagnostics = make(map[span.URI][]source.Diagnostic)
@@ -366,15 +370,26 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span) {
 			r.diagnostics[key.id.URI] = diags
 		}
 	}
-	diags := r.diagnostics[uri]
-
+	var diag *source.Diagnostic
+	for _, d := range r.diagnostics[uri] {
+		// Compare the start positions rather than the entire range because
+		// some diagnostics have a range with the same start and end position (8:1-8:1).
+		// The current marker functionality prevents us from having a range of 0 length.
+		if protocol.ComparePosition(d.Range.Start, rng.Start) == 0 {
+			diag = &d
+			break
+		}
+	}
+	if diag == nil {
+		t.Fatalf("could not get any suggested fixes for %v", spn)
+	}
 	actions, err := r.server.CodeAction(r.ctx, &protocol.CodeActionParams{
 		TextDocument: protocol.TextDocumentIdentifier{
 			URI: protocol.NewURI(uri),
 		},
 		Context: protocol.CodeActionContext{
 			Only:        []protocol.CodeActionKind{protocol.QuickFix},
-			Diagnostics: toProtocolDiagnostics(diags),
+			Diagnostics: toProtocolDiagnostics([]source.Diagnostic{*diag}),
 		},
 	})
 	if err != nil {
@@ -388,12 +403,13 @@ func (r *runner) SuggestedFix(t *testing.T, spn span.Span) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	got := res[uri]
-	fixed := string(r.data.Golden("suggestedfix", filename, func() ([]byte, error) {
-		return []byte(got), nil
-	}))
-	if fixed != got {
-		t.Errorf("suggested fixes failed for %s, expected:\n%#v\ngot:\n%#v", filename, fixed, got)
+	for u, got := range res {
+		fixed := string(r.data.Golden("suggestedfix_"+tests.SpanName(spn), u.Filename(), func() ([]byte, error) {
+			return []byte(got), nil
+		}))
+		if fixed != got {
+			t.Errorf("suggested fixes failed for %s, expected:\n%#v\ngot:\n%#v", u.Filename(), fixed, got)
+		}
 	}
 }
 
@@ -892,103 +908,3 @@ func TestBytesOffset(t *testing.T) {
 		}
 	}
 }
-
-// TODO(golang/go#36091): This function can be refactored to look like the rest of this file
-// when marker support gets added for go.mod files.
-func TestModfileSuggestedFixes(t *testing.T) {
-	t.Skip("this test cannot find example.com/package")
-
-	ctx := tests.Context(t)
-	cache := cache.New(nil)
-	session := cache.NewSession()
-	options := tests.DefaultOptions()
-	options.TempModfile = true
-	options.Env = append(os.Environ(), "GOPACKAGESDRIVER=off", "GOROOT=")
-
-	server := Server{
-		session:   session,
-		delivered: map[span.URI]sentDiagnostics{},
-	}
-
-	for _, tt := range []string{"indirect/primarymod", "unused/primarymod"} {
-		t.Run(tt, func(t *testing.T) {
-			folder, err := tests.CopyFolderToTempDir(filepath.Join("testdata", tt))
-			if err != nil {
-				t.Fatal(err)
-			}
-			defer os.RemoveAll(folder)
-
-			_, snapshot, err := session.NewView(ctx, "suggested_fix_test", span.FileURI(folder), options)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			realURI, tempURI := snapshot.View().ModFiles()
-			// TODO: Add testing for when the -modfile flag is turned off and we still get diagnostics.
-			if tempURI == "" {
-				return
-			}
-			realfh, err := snapshot.GetFile(realURI)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			reports, _, err := mod.Diagnostics(ctx, snapshot)
-			if err != nil {
-				t.Fatal(err)
-			}
-			if len(reports) != 1 {
-				t.Errorf("expected 1 fileHandle, got %d", len(reports))
-			}
-
-			_, m, _, _, err := snapshot.ModTidyHandle(ctx, realfh).Tidy(ctx)
-			if err != nil {
-				t.Fatal(err)
-			}
-
-			for fh, diags := range reports {
-				actions, err := server.CodeAction(ctx, &protocol.CodeActionParams{
-					TextDocument: protocol.TextDocumentIdentifier{
-						URI: protocol.NewURI(fh.URI),
-					},
-					Context: protocol.CodeActionContext{
-						Only:        []protocol.CodeActionKind{protocol.SourceOrganizeImports},
-						Diagnostics: toProtocolDiagnostics(diags),
-					},
-				})
-				if err != nil {
-					t.Fatal(err)
-				}
-				if len(actions) == 0 {
-					t.Fatal("no code actions returned")
-				}
-				if len(actions) > 1 {
-					t.Fatal("expected only 1 code action")
-				}
-				res := map[span.URI]string{}
-				for _, docEdits := range actions[0].Edit.DocumentChanges {
-					uri := span.NewURI(docEdits.TextDocument.URI)
-					content, err := ioutil.ReadFile(uri.Filename())
-					if err != nil {
-						t.Fatal(err)
-					}
-					res[uri] = string(content)
-					sedits, err := source.FromProtocolEdits(m, docEdits.Edits)
-					if err != nil {
-						t.Fatal(err)
-					}
-					res[uri] = applyEdits(res[uri], sedits)
-				}
-				got := res[realfh.Identity().URI]
-				contents, err := ioutil.ReadFile(filepath.Join(folder, "go.mod.golden"))
-				if err != nil {
-					t.Fatal(err)
-				}
-				want := string(contents)
-				if want != got {
-					t.Errorf("suggested fixes failed for %s, expected:\n%s\ngot:\n%s", fh.URI.Filename(), want, got)
-				}
-			}
-		})
-	}
-}

internal/lsp/mod/diagnostics.go

@@ -8,6 +8,8 @@ package mod
 
 import (
 	"context"
+	"fmt"
+	"regexp"
 
 	"golang.org/x/mod/modfile"
 	"golang.org/x/tools/internal/lsp/protocol"
@@ -108,6 +110,85 @@ func SuggestedFixes(ctx context.Context, snapshot source.Snapshot, realfh source
 	return actions
 }
 
+func SuggestedGoFixes(ctx context.Context, snapshot source.Snapshot, gofh source.FileHandle, diags []protocol.Diagnostic) ([]protocol.CodeAction, error) {
+	// TODO: We will want to support diagnostics for go.mod files even when the -modfile flag is turned off.
+	realURI, tempURI := snapshot.View().ModFiles()
+
+	// Check the case when the tempModfile flag is turned off.
+	if realURI == "" || tempURI == "" {
+		return nil, nil
+	}
+
+	ctx, done := trace.StartSpan(ctx, "mod.SuggestedGoFixes", telemetry.File.Of(realURI))
+	defer done()
+
+	realfh, err := snapshot.GetFile(realURI)
+	if err != nil {
+		return nil, err
+	}
+	realFile, realMapper, missingDeps, _, err := snapshot.ModTidyHandle(ctx, realfh).Tidy(ctx)
+	if err != nil {
+		return nil, err
+	}
+	// Get the contents of the go.mod file before we make any changes.
+	oldContents, _, err := realfh.Read(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	var actions []protocol.CodeAction
+	for _, diag := range diags {
+		re := regexp.MustCompile(`(.+) is not in your go.mod file`)
+		matches := re.FindStringSubmatch(diag.Message)
+		if len(matches) != 2 {
+			continue
+		}
+		req := missingDeps[matches[1]]
+		if req == nil {
+			continue
+		}
+		// Calculate the quick fix edits that need to be made to the go.mod file.
+		if err := realFile.AddRequire(req.Mod.Path, req.Mod.Version); err != nil {
+			return nil, err
+		}
+		realFile.Cleanup()
+		newContents, err := realFile.Format()
+		if err != nil {
+			return nil, err
+		}
+		// Reset the *modfile.File back to before we added the dependency.
+		if err := realFile.DropRequire(req.Mod.Path); err != nil {
+			return nil, err
+		}
+		// Calculate the edits to be made due to the change.
+		diff := snapshot.View().Options().ComputeEdits(realfh.Identity().URI, string(oldContents), string(newContents))
+		edits, err := source.ToProtocolEdits(realMapper, diff)
+		if err != nil {
+			return nil, err
+		}
+		action := protocol.CodeAction{
+			Title:       fmt.Sprintf("Add %s to go.mod", req.Mod.Path),
+			Kind:        protocol.QuickFix,
+			Diagnostics: []protocol.Diagnostic{diag},
+			Edit: protocol.WorkspaceEdit{
+				DocumentChanges: []protocol.TextDocumentEdit{
+					{
+						TextDocument: protocol.VersionedTextDocumentIdentifier{
+							Version: realfh.Identity().Version,
+							TextDocumentIdentifier: protocol.TextDocumentIdentifier{
+								URI: protocol.NewURI(realfh.Identity().URI),
+							},
+						},
+						Edits: edits,
+					},
+				},
+			},
+		}
+		actions = append(actions, action)
+	}
+	return actions, nil
+}
+
 func sameDiagnostic(d protocol.Diagnostic, e source.Error) bool {
 	return d.Message == e.Message && protocol.CompareRange(d.Range, e.Range) == 0 && d.Source == e.Category
 }

internal/lsp/testdata/indirect/primarymod/go.mod.golden

@@ -1,4 +1,4 @@
--- suggestedfix --
+-- suggestedfix_go.mod_5_40 --
 module indirect
 
 go 1.12

internal/lsp/testdata/lsp/primarymod/suggestedfix/has_suggested_fix.go.golden

@@ -1,4 +1,4 @@
--- suggestedfix --
+-- suggestedfix_has_suggested_fix_9_2 --
 package suggestedfix
 
 import (

internal/lsp/testdata/missingdep/modules/example.com/extramodule/pkg/x.go

@@ -0,0 +1,3 @@
+package pkg
+
+const Test = 1

internal/lsp/testdata/missingdep/primarymod/go.mod

@@ -0,0 +1,3 @@
+module missingdep
+
+go 1.12

internal/lsp/testdata/missingdep/primarymod/go.mod.golden

@@ -0,0 +1,7 @@
+-- suggestedfix_main_5_2 --
+module missingdep
+
+go 1.12
+
+require example.com/extramodule v1.0.0
+

internal/lsp/testdata/missingdep/primarymod/main.go

@@ -0,0 +1,10 @@
+// Package missingdep does something
+package missingdep
+
+import (
+	"example.com/extramodule/pkg" //@diag("\"example.com/extramodule/pkg\"", "go mod tidy", "example.com/extramodule is not in your go.mod file.", "warning"),suggestedfix("\"example.com/extramodule/pkg\"")
+)
+
+func Yo() {
+	_ = pkg.Test
+}

internal/lsp/testdata/missingdep/summary.txt.golden

@@ -0,0 +1,25 @@
+-- summary --
+CompletionsCount = 0
+CompletionSnippetCount = 0
+UnimportedCompletionsCount = 0
+DeepCompletionsCount = 0
+FuzzyCompletionsCount = 0
+RankedCompletionsCount = 0
+CaseSensitiveCompletionsCount = 0
+DiagnosticsCount = 1
+FoldingRangesCount = 0
+FormatCount = 0
+ImportCount = 0
+SuggestedFixCount = 1
+DefinitionsCount = 0
+TypeDefinitionsCount = 0
+HighlightsCount = 0
+ReferencesCount = 0
+RenamesCount = 0
+PrepareRenamesCount = 0
+SymbolsCount = 0
+WorkspaceSymbolsCount = 0
+SignaturesCount = 0
+LinksCount = 0
+ImplementationsCount = 0
+

internal/lsp/testdata/missingtwodep/modules/example.com/anothermodule/hey/y.go

@@ -0,0 +1,3 @@
+package hey
+
+const Test = 1

internal/lsp/testdata/missingtwodep/modules/example.com/extramodule/pkg/x.go

@@ -0,0 +1,3 @@
+package pkg
+
+const Test = 1

internal/lsp/testdata/missingtwodep/modules/example.com/extramodule/yo/y.go

@@ -0,0 +1,3 @@
+package yo
+
+const Test = 1

internal/lsp/testdata/missingtwodep/primarymod/go.mod

@@ -0,0 +1,3 @@
+module missingtwodep
+
+go 1.12

internal/lsp/testdata/missingtwodep/primarymod/go.mod.golden

@@ -0,0 +1,21 @@
+-- suggestedfix_main_5_2 --
+module missingtwodep
+
+go 1.12
+
+require example.com/anothermodule v1.0.0
+
+-- suggestedfix_main_6_2 --
+module missingtwodep
+
+go 1.12
+
+require example.com/extramodule v1.0.0
+
+-- suggestedfix_main_7_2 --
+module missingtwodep
+
+go 1.12
+
+require example.com/extramodule v1.0.0
+

internal/lsp/testdata/missingtwodep/primarymod/main.go

@@ -0,0 +1,14 @@
+// Package missingtwodep does something
+package missingtwodep
+
+import (
+	"example.com/anothermodule/hey" //@diag("\"example.com/anothermodule/hey\"", "go mod tidy", "example.com/anothermodule is not in your go.mod file.", "warning"),suggestedfix("\"example.com/anothermodule/hey\"")
+	"example.com/extramodule/pkg"   //@diag("\"example.com/extramodule/pkg\"", "go mod tidy", "example.com/extramodule is not in your go.mod file.", "warning"),suggestedfix("\"example.com/extramodule/pkg\"")
+	"example.com/extramodule/yo"    //@diag("\"example.com/extramodule/yo\"", "go mod tidy", "example.com/extramodule is not in your go.mod file.", "warning"),suggestedfix("\"example.com/extramodule/yo\"")
+)
+
+func Yo() {
+	_ = pkg.Test
+	_ = yo.Test
+	_ = hey.Test
+}

internal/lsp/testdata/missingtwodep/summary.txt.golden

@@ -0,0 +1,25 @@
+-- summary --
+CompletionsCount = 0
+CompletionSnippetCount = 0
+UnimportedCompletionsCount = 0
+DeepCompletionsCount = 0
+FuzzyCompletionsCount = 0
+RankedCompletionsCount = 0
+CaseSensitiveCompletionsCount = 0
+DiagnosticsCount = 3
+FoldingRangesCount = 0
+FormatCount = 0
+ImportCount = 0
+SuggestedFixCount = 3
+DefinitionsCount = 0
+TypeDefinitionsCount = 0
+HighlightsCount = 0
+ReferencesCount = 0
+RenamesCount = 0
+PrepareRenamesCount = 0
+SymbolsCount = 0
+WorkspaceSymbolsCount = 0
+SignaturesCount = 0
+LinksCount = 0
+ImplementationsCount = 0
+

internal/lsp/testdata/unused/primarymod/go.mod.golden

@@ -1,4 +1,4 @@
--- suggestedfix --
+-- suggestedfix_go.mod_5_1 --
 module unused
 
 go 1.12

internal/lsp/tests/tests.go

@@ -416,7 +416,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 
 		for src, exp := range cases {
 			for i, e := range exp {
-				t.Run(spanName(src)+"_"+strconv.Itoa(i), func(t *testing.T) {
+				t.Run(SpanName(src)+"_"+strconv.Itoa(i), func(t *testing.T) {
 					t.Helper()
 					if (!haveCgo || runtime.GOOS == "android") && strings.Contains(t.Name(), "cgo") {
 						t.Skip("test requires cgo, not supported")
@@ -438,7 +438,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 		for _, placeholders := range []bool{true, false} {
 			for src, expecteds := range data.CompletionSnippets {
 				for i, expected := range expecteds {
-					name := spanName(src) + "_" + strconv.Itoa(i+1)
+					name := SpanName(src) + "_" + strconv.Itoa(i+1)
 					if placeholders {
 						name += "_placeholders"
 					}
@@ -480,9 +480,8 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("Diagnostics", func(t *testing.T) {
 		t.Helper()
 		for uri, want := range data.Diagnostics {
-			// If the -modfile flag is not available, then we do not want to run
-			// diagnostics on any .mod file.
-			if strings.Contains(uri.Filename(), ".mod") && !data.ModfileFlagAvailable {
+			// Check if we should skip this URI if the -modfile flag is not available.
+			if shouldSkip(data, uri) {
 				continue
 			}
 			t.Run(uriName(uri), func(t *testing.T) {
@@ -525,12 +524,11 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("SuggestedFix", func(t *testing.T) {
 		t.Helper()
 		for _, spn := range data.SuggestedFixes {
-			// If the -modfile flag is not available, then we do not want to run
-			// suggested fixes on any .mod file.
-			if strings.Contains(spn.URI().Filename(), ".mod") && !data.ModfileFlagAvailable {
+			// Check if we should skip this spn if the -modfile flag is not available.
+			if shouldSkip(data, spn.URI()) {
 				continue
 			}
-			t.Run(spanName(spn), func(t *testing.T) {
+			t.Run(SpanName(spn), func(t *testing.T) {
 				t.Helper()
 				tests.SuggestedFix(t, spn)
 			})
@@ -540,7 +538,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("Definition", func(t *testing.T) {
 		t.Helper()
 		for spn, d := range data.Definitions {
-			t.Run(spanName(spn), func(t *testing.T) {
+			t.Run(SpanName(spn), func(t *testing.T) {
 				t.Helper()
 				if (!haveCgo || runtime.GOOS == "android") && strings.Contains(t.Name(), "cgo") {
 					t.Skip("test requires cgo, not supported")
@@ -553,7 +551,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("Implementation", func(t *testing.T) {
 		t.Helper()
 		for spn, m := range data.Implementations {
-			t.Run(spanName(spn), func(t *testing.T) {
+			t.Run(SpanName(spn), func(t *testing.T) {
 				t.Helper()
 				tests.Implementation(t, spn, m)
 			})
@@ -563,7 +561,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("Highlight", func(t *testing.T) {
 		t.Helper()
 		for pos, locations := range data.Highlights {
-			t.Run(spanName(pos), func(t *testing.T) {
+			t.Run(SpanName(pos), func(t *testing.T) {
 				t.Helper()
 				tests.Highlight(t, pos, locations)
 			})
@@ -573,7 +571,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("References", func(t *testing.T) {
 		t.Helper()
 		for src, itemList := range data.References {
-			t.Run(spanName(src), func(t *testing.T) {
+			t.Run(SpanName(src), func(t *testing.T) {
 				t.Helper()
 				tests.References(t, src, itemList)
 			})
@@ -593,7 +591,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("PrepareRenames", func(t *testing.T) {
 		t.Helper()
 		for src, want := range data.PrepareRenames {
-			t.Run(spanName(src), func(t *testing.T) {
+			t.Run(SpanName(src), func(t *testing.T) {
 				t.Helper()
 				tests.PrepareRename(t, src, want)
 			})
@@ -634,7 +632,7 @@ func Run(t *testing.T, tests Tests, data *Data) {
 	t.Run("SignatureHelp", func(t *testing.T) {
 		t.Helper()
 		for spn, expectedSignature := range data.Signatures {
-			t.Run(spanName(spn), func(t *testing.T) {
+			t.Run(SpanName(spn), func(t *testing.T) {
 				t.Helper()
 				tests.SignatureHelp(t, spn, expectedSignature)
 			})
@@ -1043,7 +1041,7 @@ func uriName(uri span.URI) string {
 	return filepath.Base(strings.TrimSuffix(uri.Filename(), ".go"))
 }
 
-func spanName(spn span.Span) string {
+func SpanName(spn span.Span) string {
 	return fmt.Sprintf("%v_%v_%v", uriName(spn.URI()), spn.Start().Line(), spn.Start().Column())
 }
 
@@ -1100,3 +1098,18 @@ func testFolders(root string) ([]string, error) {
 	}
 	return folders, nil
 }
+
+func shouldSkip(data *Data, uri span.URI) bool {
+	if data.ModfileFlagAvailable {
+		return false
+	}
+	// If the -modfile flag is not available, then we do not want to run
+	// any tests on the go.mod file.
+	if strings.Contains(uri.Filename(), ".mod") {
+		return true
+	}
+	// If the -modfile flag is not available, then we do not want to test any
+	// uri that contains "go mod tidy".
+	m, err := data.Mapper(uri)
+	return err == nil && strings.Contains(string(m.Content), ", \"go mod tidy\",")
+}