gopls/internal/golang: implement Go to TestXxx CodeLens

DO NOT REVIEW
DO NOT SUBMIT

Author: ShoshinNikita

gopls/internal/cache/snapshot.go

@@ -976,7 +976,7 @@ func (s *Snapshot) WorkspaceMetadata(ctx context.Context) ([]*metadata.Package,
 	defer s.mu.Unlock()
 
 	meta := make([]*metadata.Package, 0, s.workspacePackages.Len())
-	for id, _ := range s.workspacePackages.All() {
+	for id := range s.workspacePackages.All() {
 		meta = append(meta, s.meta.Packages[id])
 	}
 	return meta, nil

gopls/internal/golang/code_lens.go

@@ -5,19 +5,25 @@
 package golang
 
 import (
+	"cmp"
 	"context"
+	"fmt"
 	"go/ast"
 	"go/token"
 	"go/types"
 	"regexp"
+	"slices"
 	"strings"
+	"unicode"
 
 	"golang.org/x/tools/gopls/internal/cache"
+	"golang.org/x/tools/gopls/internal/cache/metadata"
 	"golang.org/x/tools/gopls/internal/cache/parsego"
 	"golang.org/x/tools/gopls/internal/file"
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/protocol/command"
 	"golang.org/x/tools/gopls/internal/settings"
+	"golang.org/x/tools/gopls/internal/util/astutil"
 )
 
 // CodeLensSources returns the supported sources of code lenses for Go files.
@@ -26,6 +32,7 @@ func CodeLensSources() map[settings.CodeLensSource]cache.CodeLensSourceFunc {
 		settings.CodeLensGenerate:      goGenerateCodeLens, // commands: Generate
 		settings.CodeLensTest:          runTestCodeLens,    // commands: Test
 		settings.CodeLensRegenerateCgo: regenerateCgoLens,  // commands: RegenerateCgo
+		settings.CodeLensGoToTest:      goToTestCodeLens,   // commands: GoToTest
 	}
 }
 
@@ -204,3 +211,143 @@ func regenerateCgoLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Ha
 	cmd := command.NewRegenerateCgoCommand("regenerate cgo definitions", command.URIArg{URI: puri})
 	return []protocol.CodeLens{{Range: rng, Command: cmd}}, nil
 }
+
+func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle) ([]protocol.CodeLens, error) {
+	if strings.HasSuffix(fh.URI().Path(), "_test.go") {
+		// Ignore test files.
+		return nil, nil
+	}
+
+	// Inspect all packages to cover both "p [p.test]" and "p_test [p.test]".
+	allPackages, err := snapshot.WorkspaceMetadata(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't request workspace metadata: %w", err)
+	}
+	dir := fh.URI().Dir()
+	testPackages := slices.DeleteFunc(allPackages, func(meta *metadata.Package) bool {
+		if meta.IsIntermediateTestVariant() || len(meta.CompiledGoFiles) == 0 || meta.ForTest == "" {
+			return true
+		}
+		return meta.CompiledGoFiles[0].Dir() != dir
+	})
+	if len(testPackages) == 0 {
+		return nil, nil
+	}
+
+	pgf, err := snapshot.ParseGo(ctx, fh, parsego.Full)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't parse file: %w", err)
+	}
+	funcPos := make(map[string]protocol.Position)
+	for _, d := range pgf.File.Decls {
+		fn, ok := d.(*ast.FuncDecl)
+		if !ok {
+			continue
+		}
+		rng, err := pgf.NodeRange(fn)
+		if err != nil {
+			return nil, fmt.Errorf("couldn't get node range: %w", err)
+		}
+
+		name := fn.Name.Name
+		if fn.Recv != nil && len(fn.Recv.List) > 0 {
+			_, rname, _ := astutil.UnpackRecv(fn.Recv.List[0].Type)
+			name = rname.Name + "_" + fn.Name.Name
+		}
+		funcPos[name] = rng.Start
+	}
+
+	type TestType int
+
+	// Types are sorted by priority from high to low.
+	const (
+		T TestType = iota + 1
+		E
+		B
+		F
+	)
+	testTypes := map[string]TestType{
+		"Test":      T,
+		"Example":   E,
+		"Benchmark": B,
+		"Fuzz":      F,
+	}
+
+	type Test struct {
+		FuncPos protocol.Position
+		Name    string
+		Loc     protocol.Location
+		Type    TestType
+	}
+	var matchedTests []Test
+
+	pkgIDs := make([]PackageID, 0, len(testPackages))
+	for _, pkg := range testPackages {
+		pkgIDs = append(pkgIDs, pkg.ID)
+	}
+	allTests, err := snapshot.Tests(ctx, pkgIDs...)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't request all tests for packages %v: %w", pkgIDs, err)
+	}
+	for _, tests := range allTests {
+		for _, test := range tests.All() {
+			var (
+				name     string
+				testType TestType
+			)
+			for prefix, t := range testTypes {
+				if strings.HasPrefix(test.Name, prefix) {
+					testType = t
+					name = test.Name[len(prefix):]
+					break
+				}
+			}
+			if testType == 0 {
+				continue // unknown type
+			}
+			name = strings.TrimPrefix(name, "_")
+
+			// Try to find 'Foo' for 'TestFoo' and 'foo' for 'Test_foo'.
+			pos, ok := funcPos[name]
+			if !ok && token.IsExported(name) {
+				// Try to find 'foo' for 'TestFoo'.
+				runes := []rune(name)
+				runes[0] = unicode.ToLower(runes[0])
+				pos, ok = funcPos[string(runes)]
+			}
+			if ok {
+				loc := test.Location
+				loc.Range.End = loc.Range.Start // move cursor to the test's beginning
+
+				matchedTests = append(matchedTests, Test{
+					FuncPos: pos,
+					Name:    test.Name,
+					Loc:     loc,
+					Type:    testType,
+				})
+			}
+		}
+	}
+	if len(matchedTests) == 0 {
+		return nil, nil
+	}
+
+	slices.SortFunc(matchedTests, func(a, b Test) int {
+		if v := protocol.ComparePosition(a.FuncPos, b.FuncPos); v != 0 {
+			return v
+		}
+		if v := cmp.Compare(a.Type, b.Type); v != 0 {
+			return v
+		}
+		return cmp.Compare(a.Name, b.Name)
+	})
+
+	lenses := make([]protocol.CodeLens, 0, len(matchedTests))
+	for _, t := range matchedTests {
+		lenses = append(lenses, protocol.CodeLens{
+			Range:   protocol.Range{Start: t.FuncPos, End: t.FuncPos},
+			Command: command.NewGoToTestCommand("Go to "+t.Name, t.Loc),
+		})
+	}
+	return lenses, nil
+}

gopls/internal/protocol/command/command_gen.go

@@ -305,6 +305,9 @@ type Interface interface {
 
 	// ModifyTags: Add or remove struct tags on a given node.
 	ModifyTags(context.Context, ModifyTagsArgs) error
+
+	// GoToTest: Go to test declaration.
+	GoToTest(context.Context, protocol.Location) error
 }
 
 type RunTestsArgs struct {

gopls/internal/protocol/command/interface.go

@@ -1788,6 +1788,15 @@ func optionsStringToMap(options string) (map[string][]string, error) {
 	return optionsMap, nil
 }
 
+func (c *commandHandler) GoToTest(ctx context.Context, loc protocol.Location) error {
+	return c.run(ctx, commandConfig{
+		forURI: loc.URI,
+	}, func(ctx context.Context, deps commandDeps) error {
+		showDocumentImpl(ctx, c.s.client, protocol.URI(loc.URI), &loc.Range, c.s.options)
+		return nil
+	})
+}
+
 func (c *commandHandler) ModifyTags(ctx context.Context, args command.ModifyTagsArgs) error {
 	return c.run(ctx, commandConfig{
 		progress: "Modifying tags",

gopls/internal/server/command.go

@@ -356,6 +356,12 @@ const (
 	// module root so that it contains an up-to-date copy of all
 	// necessary package dependencies.
 	CodeLensVendor CodeLensSource = "vendor"
+
+	// Go to the functions's Test, Example, Benchmark, or Fuzz declarations
+	//
+	// This codelens source annotates the function and method declarations
+	// with their corresponding Test, Example, Benchmark, and Fuzz functions.
+	CodeLensGoToTest CodeLensSource = "go_to_test"
 )
 
 // Note: CompletionOptions must be comparable with reflect.DeepEqual.

gopls/internal/settings/settings.go

@@ -0,0 +1,62 @@
+-- settings.json --
+{
+	"codelenses": {
+		"go_to_test": true
+	}
+}
+
+-- a.go --
+//@codelenses()
+
+package codelenses
+
+func Foo() int { return 0 } //@codelens(re"()func Foo", "Go to TestFoo")
+
+func bar() {} //@codelens(re"()func bar", "Go to Test_bar"),codelens(re"()func bar", "Go to Benchmark_bar"),codelens(re"()func bar", "Go to Fuzz_bar")
+
+func xyz() {} //@codelens(re"()func xyz", "Go to TestXyz")
+
+-- a1_test.go --
+package codelenses
+
+import "testing"
+
+func TestFoo(*testing.T) {}
+
+-- a2_test.go --
+package codelenses
+
+import "testing"
+
+func TestXyz(*testing.T) {}
+
+-- a3_test.go --
+package codelenses_test
+
+import "testing"
+
+func Test_bar(*testing.T) {}
+
+func Benchmark_bar(*testing.B) {}
+
+func Fuzz_bar(*testing.F) {}
+
+-- writer.go --
+//@codelenses()
+
+package codelenses
+
+type Writer struct{}
+
+func (w *Writer) Write() error { return nil } //@codelens(re"()func", "Go to TestWriter_Write")
+
+func (w *Writer) Flush() {} //@codelens(re"()func", "Go to TestWriter_Flush")
+
+-- writer_test.go --
+package codelenses
+
+import "testing"
+
+func TestWriter_Write(*testing.T) {}
+
+func TestWriter_Flush(*testing.T) {}