internal/lsp: clean up some of the mod code lens code

Refactor the checks for code lenses being enabled out of the source
package so that the mod code lenses can also make use of them.

Also, a few small changes to the titles of the `go mod tidy` and `go mod
vendor` code lenses.

Change-Id: I4e79ab08a4e7aea4d4d6de6fd652d0b77d30c811
Reviewed-on: https://go-review.googlesource.com/c/tools/+/252397
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
diff --git a/internal/lsp/code_lens.go b/internal/lsp/code_lens.go
index 75b6301..558344d 100644
--- a/internal/lsp/code_lens.go
+++ b/internal/lsp/code_lens.go
@@ -6,7 +6,9 @@
 
 import (
 	"context"
+	"fmt"
 
+	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/lsp/mod"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
@@ -18,12 +20,29 @@
 	if !ok {
 		return nil, err
 	}
+	var lensFuncs map[string]source.LensFunc
 	switch fh.Kind() {
 	case source.Mod:
-		return mod.CodeLens(ctx, snapshot, fh.URI())
+		lensFuncs = mod.LensFuncs()
 	case source.Go:
-		return source.CodeLens(ctx, snapshot, fh)
+		lensFuncs = source.LensFuncs()
+	default:
+		// Unsupported file kind for a code lens.
+		return nil, nil
 	}
-	// Unsupported file kind for a code action.
-	return nil, nil
+	var result []protocol.CodeLens
+	for lens, lf := range lensFuncs {
+		if !snapshot.View().Options().EnabledCodeLens[lens] {
+			continue
+		}
+		added, err := lf(ctx, snapshot, fh)
+		// Code lens is called on every keystroke, so we should just operate in
+		// a best-effort mode, ignoring errors.
+		if err != nil {
+			event.Error(ctx, fmt.Sprintf("code lens %s failed", lens), err)
+			continue
+		}
+		result = append(result, added...)
+	}
+	return result, nil
 }
diff --git a/internal/lsp/lsp_test.go b/internal/lsp/lsp_test.go
index c287276..20459f1 100644
--- a/internal/lsp/lsp_test.go
+++ b/internal/lsp/lsp_test.go
@@ -18,7 +18,6 @@
 	"golang.org/x/tools/internal/lsp/cache"
 	"golang.org/x/tools/internal/lsp/diff"
 	"golang.org/x/tools/internal/lsp/diff/myers"
-	"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"
@@ -163,18 +162,16 @@
 	if source.DetectLanguage("", uri.Filename()) != source.Mod {
 		return
 	}
-	v, err := r.server.session.ViewOf(uri)
-	if err != nil {
-		t.Fatal(err)
-	}
-	snapshot, release := v.Snapshot(r.ctx)
-	defer release()
-	got, err := mod.CodeLens(r.ctx, snapshot, uri)
+	got, err := r.server.codeLens(r.ctx, &protocol.CodeLensParams{
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: protocol.DocumentURI(uri),
+		},
+	})
 	if err != nil {
 		t.Fatal(err)
 	}
 	if diff := tests.DiffCodeLens(uri, want, got); diff != "" {
-		t.Error(diff)
+		t.Errorf("%s: %s", uri, diff)
 	}
 }
 
diff --git a/internal/lsp/mod/code_lens.go b/internal/lsp/mod/code_lens.go
index 6e7b707..4f926fa 100644
--- a/internal/lsp/mod/code_lens.go
+++ b/internal/lsp/mod/code_lens.go
@@ -3,45 +3,41 @@
 import (
 	"context"
 	"fmt"
+	"os"
+	"path/filepath"
 
 	"golang.org/x/mod/modfile"
-	"golang.org/x/tools/internal/event"
-	"golang.org/x/tools/internal/lsp/debug/tag"
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/span"
 )
 
-// CodeLens computes code lens for a go.mod file.
-func CodeLens(ctx context.Context, snapshot source.Snapshot, uri span.URI) ([]protocol.CodeLens, error) {
-	if !snapshot.View().Options().EnabledCodeLens[source.CommandUpgradeDependency.Name] {
-		return nil, nil
+// LensFuncs returns the supported lensFuncs for go.mod files.
+func LensFuncs() map[string]source.LensFunc {
+	return map[string]source.LensFunc{
+		source.CommandUpgradeDependency.Name: upgradeLens,
+		source.CommandTidy.Name:              tidyLens,
+		source.CommandVendor.Name:            vendorLens,
 	}
-	ctx, done := event.Start(ctx, "mod.CodeLens", tag.URI.Of(uri))
-	defer done()
+}
 
-	// Only show go.mod code lenses in module mode, for the view's go.mod.
-	if modURI := snapshot.View().ModFile(); modURI == "" || modURI != uri {
-		return nil, nil
-	}
-	fh, err := snapshot.GetFile(ctx, uri)
-	if err != nil {
-		return nil, err
-	}
+func upgradeLens(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.CodeLens, error) {
 	pm, err := snapshot.ParseMod(ctx, fh)
 	if err != nil {
 		return nil, err
 	}
+	module := pm.File.Module
+	if module == nil || module.Syntax == nil {
+		return nil, nil
+	}
 	upgrades, err := snapshot.ModUpgrade(ctx, fh)
 	if err != nil {
 		return nil, err
 	}
-
 	var (
-		codelens    []protocol.CodeLens
+		codelenses  []protocol.CodeLens
 		allUpgrades []string
 	)
-
 	for _, req := range pm.File.Require {
 		dep := req.Mod.Path
 		latest, ok := upgrades[dep]
@@ -49,15 +45,15 @@
 			continue
 		}
 		// Get the range of the require directive.
-		rng, err := positionsToRange(uri, pm.Mapper, req.Syntax.Start, req.Syntax.End)
+		rng, err := positionsToRange(fh.URI(), pm.Mapper, req.Syntax.Start, req.Syntax.End)
 		if err != nil {
 			return nil, err
 		}
-		upgradeDepArgs, err := source.MarshalArgs(uri, []string{dep})
+		upgradeDepArgs, err := source.MarshalArgs(fh.URI(), []string{dep})
 		if err != nil {
 			return nil, err
 		}
-		codelens = append(codelens, protocol.CodeLens{
+		codelenses = append(codelenses, protocol.CodeLens{
 			Range: rng,
 			Command: protocol.Command{
 				Title:     fmt.Sprintf("Upgrade dependency to %s", latest),
@@ -67,53 +63,20 @@
 		})
 		allUpgrades = append(allUpgrades, dep)
 	}
-
-	module := pm.File.Module
-	if module == nil || module.Syntax == nil {
-		return codelens, nil
-	}
-	// Get the range of the module directive.
-	rng, err := positionsToRange(uri, pm.Mapper, module.Syntax.Start, module.Syntax.End)
-	if err != nil {
-		return nil, err
-	}
-
-	// Add go mod vendor and go mod tidy lenses
-	goModArgs, err := source.MarshalArgs(uri)
-	if err != nil {
-		return nil, err
-	}
-	codelens = append(codelens, protocol.CodeLens{
-		Range: rng,
-		Command: protocol.Command{
-			Title:     "Sync vendor directory",
-			Command:   source.CommandVendor.Name,
-			Arguments: goModArgs,
-		},
-	})
-	tidied, err := snapshot.ModTidy(ctx, fh)
-	if err != nil && err != source.ErrTmpModfileUnsupported {
-		return nil, err
-	}
-	if tidied != nil && len(tidied.Errors) > 0 {
-		codelens = append(codelens, protocol.CodeLens{
-			Range: rng,
-			Command: protocol.Command{
-				Title:     "Remove unused dependencies",
-				Command:   source.CommandTidy.Name,
-				Arguments: goModArgs,
-			},
-		})
-	}
-
-	// If there is at least 1 upgrade, add an "Upgrade all dependencies" to the module statement.
+	// If there is at least 1 upgrade, add "Upgrade all dependencies" to
+	// the module statement.
 	if len(allUpgrades) > 0 {
-		upgradeDepArgs, err := source.MarshalArgs(uri, append([]string{"-u"}, allUpgrades...))
+		upgradeDepArgs, err := source.MarshalArgs(fh.URI(), append([]string{"-u"}, allUpgrades...))
 		if err != nil {
 			return nil, err
 		}
-		codelens = append(codelens, protocol.CodeLens{
-			Range: rng,
+		// Get the range of the module directive.
+		moduleRng, err := positionsToRange(pm.Mapper.URI, pm.Mapper, module.Syntax.Start, module.Syntax.End)
+		if err != nil {
+			return nil, err
+		}
+		codelenses = append(codelenses, protocol.CodeLens{
+			Range: moduleRng,
 			Command: protocol.Command{
 				Title:     "Upgrade all dependencies",
 				Command:   source.CommandUpgradeDependency.Name,
@@ -121,7 +84,67 @@
 			},
 		})
 	}
-	return codelens, nil
+	return codelenses, err
+}
+
+func tidyLens(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.CodeLens, error) {
+	goModArgs, err := source.MarshalArgs(fh.URI())
+	if err != nil {
+		return nil, err
+	}
+	tidied, err := snapshot.ModTidy(ctx, fh)
+	if err != nil {
+		return nil, err
+	}
+	if len(tidied.Errors) == 0 {
+		return nil, nil
+	}
+	pm, err := snapshot.ParseMod(ctx, fh)
+	if err != nil {
+		return nil, err
+	}
+	rng, err := positionsToRange(pm.Mapper.URI, pm.Mapper, pm.File.Module.Syntax.Start, pm.File.Module.Syntax.End)
+	if err != nil {
+		return nil, err
+	}
+	return []protocol.CodeLens{{
+		Range: rng,
+		Command: protocol.Command{
+			Title:     "Tidy module",
+			Command:   source.CommandTidy.Name,
+			Arguments: goModArgs,
+		},
+	}}, err
+}
+
+func vendorLens(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.CodeLens, error) {
+	goModArgs, err := source.MarshalArgs(fh.URI())
+	if err != nil {
+		return nil, err
+	}
+	pm, err := snapshot.ParseMod(ctx, fh)
+	if err != nil {
+		return nil, err
+	}
+	rng, err := positionsToRange(pm.Mapper.URI, pm.Mapper, pm.File.Module.Syntax.Start, pm.File.Module.Syntax.End)
+	if err != nil {
+		return nil, err
+	}
+	// Change the message depending on whether or not the module already has a
+	// vendor directory.
+	title := "Create vendor directory"
+	vendorDir := filepath.Join(filepath.Dir(fh.URI().Filename()), "vendor")
+	if info, _ := os.Stat(vendorDir); info != nil && info.IsDir() {
+		title = "Sync vendor directory"
+	}
+	return []protocol.CodeLens{{
+		Range: rng,
+		Command: protocol.Command{
+			Title:     title,
+			Command:   source.CommandVendor.Name,
+			Arguments: goModArgs,
+		},
+	}}, nil
 }
 
 func positionsToRange(uri span.URI, m *protocol.ColumnMapper, s, e modfile.Position) (protocol.Range, error) {
diff --git a/internal/lsp/regtest/codelens_test.go b/internal/lsp/regtest/codelens_test.go
index c6bfc10..e4c90bd 100644
--- a/internal/lsp/regtest/codelens_test.go
+++ b/internal/lsp/regtest/codelens_test.go
@@ -174,7 +174,7 @@
 	runner.Run(t, shouldRemoveDep, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
 		lenses := env.CodeLens("go.mod")
-		want := "Remove unused dependencies"
+		want := "Tidy module"
 		var found protocol.CodeLens
 		for _, lens := range lenses {
 			if lens.Command.Title == want {
@@ -183,7 +183,7 @@
 			}
 		}
 		if found.Command.Command == "" {
-			t.Fatalf("did not find lens %q, got %v", want, lenses)
+			t.Fatalf("did not find lens %q, got %v", want, found.Command)
 		}
 		if _, err := env.Editor.Server.ExecuteCommand(env.Ctx, &protocol.ExecuteCommandParams{
 			Command:   found.Command.Command,
diff --git a/internal/lsp/source/code_lens.go b/internal/lsp/source/code_lens.go
index 149af3e..966dcf2 100644
--- a/internal/lsp/source/code_lens.go
+++ b/internal/lsp/source/code_lens.go
@@ -17,29 +17,16 @@
 	"golang.org/x/tools/internal/span"
 )
 
-type lensFunc func(context.Context, Snapshot, FileHandle) ([]protocol.CodeLens, error)
+type LensFunc func(context.Context, Snapshot, FileHandle) ([]protocol.CodeLens, error)
 
-var lensFuncs = map[string]lensFunc{
-	CommandGenerate.Name:      goGenerateCodeLens,
-	CommandTest.Name:          runTestCodeLens,
-	CommandRegenerateCgo.Name: regenerateCgoLens,
-	CommandToggleDetails.Name: toggleDetailsCodeLens,
-}
-
-// CodeLens computes code lens for Go source code.
-func CodeLens(ctx context.Context, snapshot Snapshot, fh FileHandle) ([]protocol.CodeLens, error) {
-	var result []protocol.CodeLens
-	for lens, lf := range lensFuncs {
-		if !snapshot.View().Options().EnabledCodeLens[lens] {
-			continue
-		}
-		added, err := lf(ctx, snapshot, fh)
-		if err != nil {
-			return nil, err
-		}
-		result = append(result, added...)
+// LensFuncs returns the supported lensFuncs for Go files.
+func LensFuncs() map[string]LensFunc {
+	return map[string]LensFunc{
+		CommandGenerate.Name:      goGenerateCodeLens,
+		CommandTest.Name:          runTestCodeLens,
+		CommandRegenerateCgo.Name: regenerateCgoLens,
+		CommandToggleDetails.Name: toggleDetailsCodeLens,
 	}
-	return result, nil
 }
 
 var (
diff --git a/internal/lsp/source/options.go b/internal/lsp/source/options.go
index 06d4ea0..b9adbb9 100644
--- a/internal/lsp/source/options.go
+++ b/internal/lsp/source/options.go
@@ -99,9 +99,11 @@
 			CompletionDocumentation: true,
 			EnabledCodeLens: map[string]bool{
 				CommandGenerate.Name:          true,
-				CommandUpgradeDependency.Name: true,
 				CommandRegenerateCgo.Name:     true,
+				CommandTidy.Name:              true,
 				CommandToggleDetails.Name:     false,
+				CommandUpgradeDependency.Name: true,
+				CommandVendor.Name:            true,
 			},
 			ExpandWorkspaceToModule: true,
 		},
diff --git a/internal/lsp/source/source_test.go b/internal/lsp/source/source_test.go
index 85f55bd..c29899a 100644
--- a/internal/lsp/source/source_test.go
+++ b/internal/lsp/source/source_test.go
@@ -924,20 +924,7 @@
 func (r *runner) Link(t *testing.T, uri span.URI, wantLinks []tests.Link)         {}
 func (r *runner) SuggestedFix(t *testing.T, spn span.Span, actionKinds []string)  {}
 func (r *runner) FunctionExtraction(t *testing.T, start span.Span, end span.Span) {}
-
-func (r *runner) CodeLens(t *testing.T, uri span.URI, want []protocol.CodeLens) {
-	fh, err := r.snapshot.GetFile(r.ctx, uri)
-	if err != nil {
-		t.Fatal(err)
-	}
-	got, err := source.CodeLens(r.ctx, r.snapshot, fh)
-	if err != nil {
-		t.Fatal(err)
-	}
-	if diff := tests.DiffCodeLens(uri, want, got); diff != "" {
-		t.Error(diff)
-	}
-}
+func (r *runner) CodeLens(t *testing.T, uri span.URI, want []protocol.CodeLens)   {}
 
 func spanToRange(data *tests.Data, spn span.Span) (*protocol.ColumnMapper, protocol.Range, error) {
 	m, err := data.Mapper(spn.URI())
diff --git a/internal/lsp/testdata/upgradedep/primarymod/go.mod b/internal/lsp/testdata/upgradedep/primarymod/go.mod
index 911b589..dca97be 100644
--- a/internal/lsp/testdata/upgradedep/primarymod/go.mod
+++ b/internal/lsp/testdata/upgradedep/primarymod/go.mod
@@ -1,4 +1,4 @@
-module upgradedep //@codelens("module upgradedep", "Sync vendor directory", "vendor"),codelens("module upgradedep", "Upgrade all dependencies", "upgrade_dependency")
+module upgradedep //@codelens("module upgradedep", "Create vendor directory", "vendor"),codelens("module upgradedep", "Upgrade all dependencies", "upgrade_dependency")
 
 // TODO(microsoft/vscode-go#12): Another issue. //@link(`microsoft/vscode-go#12`, `https://github.com/microsoft/vscode-go/issues/12`)
 
diff --git a/internal/lsp/tests/util.go b/internal/lsp/tests/util.go
index 5d95d00..c3633a1 100644
--- a/internal/lsp/tests/util.go
+++ b/internal/lsp/tests/util.go
@@ -227,12 +227,12 @@
 	}
 	for i, w := range want {
 		g := got[i]
+		if w.Command.Command != g.Command.Command {
+			return summarizeCodeLens(i, uri, want, got, "incorrect Command Name got %v want %v", g.Command.Command, w.Command.Command)
+		}
 		if w.Command.Title != g.Command.Title {
 			return summarizeCodeLens(i, uri, want, got, "incorrect Command Title got %v want %v", g.Command.Title, w.Command.Title)
 		}
-		if w.Command.Command != g.Command.Command {
-			return summarizeCodeLens(i, uri, want, got, "incorrect Command Title got %v want %v", g.Command.Command, w.Command.Command)
-		}
 		if protocol.ComparePosition(w.Range.Start, g.Range.Start) != 0 {
 			return summarizeCodeLens(i, uri, want, got, "incorrect Start got %v want %v", g.Range.Start, w.Range.Start)
 		}
@@ -252,11 +252,11 @@
 		}
 		if c[i].Command.Command < c[j].Command.Command {
 			return true
+		} else if c[i].Command.Command == c[j].Command.Command {
+			return c[i].Command.Title < c[j].Command.Title
+		} else {
+			return false
 		}
-		if c[i].Command.Command == c[j].Command.Command {
-			return true
-		}
-		return c[i].Command.Title <= c[j].Command.Title
 	})
 }