prefix match + code action

Author: ShoshinNikita

gopls/internal/cache/testfuncs/tests.go

@@ -49,8 +49,20 @@ func (index *Index) All() []Result {
 type Result struct {
 	Location protocol.Location // location of the test
 	Name     string            // name of the test
+	Type     TestType          // type of the test
+	Subtest  bool
 }
 
+type TestType int
+
+const (
+	TypeInvalid TestType = iota
+	TypeTest
+	TypeBenchmark
+	TypeFuzz
+	TypeExample
+)
+
 // NewIndex returns a new index of method-set information for all
 // package-level types in the specified package.
 func NewIndex(files []*parsego.File, info *types.Info) *Index {
@@ -84,15 +96,16 @@ func (b *indexBuilder) build(files []*parsego.File, info *types.Info) *Index {
 				continue
 			}
 
-			isTest, isExample := isTestOrExample(obj)
-			if !isTest && !isExample {
+			testType := getTestType(obj)
+			if testType == TypeInvalid {
 				continue
 			}
 
 			var t gobTest
 			t.Name = decl.Name.Name
 			t.Location.URI = file.URI
 			t.Location.Range, _ = file.NodeRange(decl)
+			t.Type = testType
 
 			i, ok := b.fileIndex[t.Location.URI]
 			if !ok {
@@ -105,7 +118,8 @@ func (b *indexBuilder) build(files []*parsego.File, info *types.Info) *Index {
 			b.visited[obj] = true
 
 			// Check for subtests
-			if isTest {
+			switch testType {
+			case TypeTest, TypeBenchmark, TypeFuzz:
 				b.Files[i].Tests = append(b.Files[i].Tests, b.findSubtests(t, decl.Type, decl.Body, file, files, info)...)
 			}
 		}
@@ -168,6 +182,8 @@ func (b *indexBuilder) findSubtests(parent gobTest, typ *ast.FuncType, body *ast
 		t.Name = b.uniqueName(parent.Name, rewrite(constant.StringVal(val)))
 		t.Location.URI = file.URI
 		t.Location.Range, _ = file.NodeRange(call)
+		t.Type = parent.Type
+		t.Subtest = true
 		tests = append(tests, t)
 
 		fn, typ, body := findFunc(files, info, body, call.Args[1])
@@ -182,7 +198,8 @@ func (b *indexBuilder) findSubtests(parent gobTest, typ *ast.FuncType, body *ast
 		}
 
 		// Never recurse if the second argument is a top-level test function
-		if isTest, _ := isTestOrExample(fn); isTest {
+		switch getTestType(fn) {
+		case TypeTest, TypeBenchmark, TypeFuzz:
 			continue
 		}
 
@@ -258,30 +275,35 @@ func findFunc(files []*parsego.File, info *types.Info, body *ast.BlockStmt, expr
 	return nil, nil, nil
 }
 
-// isTestOrExample reports whether the given func is a testing func or an
-// example func (or neither). isTestOrExample returns (true, false) for testing
-// funcs, (false, true) for example funcs, and (false, false) otherwise.
-func isTestOrExample(fn *types.Func) (isTest, isExample bool) {
+// getTestType reports the test type of the given function.
+func getTestType(fn *types.Func) TestType {
 	sig := fn.Type().(*types.Signature)
-	if sig.Params().Len() == 0 &&
-		sig.Results().Len() == 0 {
-		return false, isTestName(fn.Name(), "Example")
+	if sig.Params().Len() == 0 && sig.Results().Len() == 0 {
+		if isTestName(fn.Name(), "Example") {
+			return TypeExample
+		}
+		return TypeInvalid
 	}
 
 	kind, ok := testKind(sig)
 	if !ok {
-		return false, false
+		return TypeInvalid
 	}
 	switch kind.Name() {
 	case "T":
-		return isTestName(fn.Name(), "Test"), false
+		if isTestName(fn.Name(), "Test") {
+			return TypeTest
+		}
 	case "B":
-		return isTestName(fn.Name(), "Benchmark"), false
+		if isTestName(fn.Name(), "Benchmark") {
+			return TypeBenchmark
+		}
 	case "F":
-		return isTestName(fn.Name(), "Fuzz"), false
-	default:
-		return false, false // "can't happen" (see testKind)
+		if isTestName(fn.Name(), "Fuzz") {
+			return TypeFuzz
+		}
 	}
+	return TypeInvalid
 }
 
 // isTestName reports whether name is a valid test name for the test kind
@@ -352,6 +374,8 @@ type gobFile struct {
 type gobTest struct {
 	Location protocol.Location // location of the test
 	Name     string            // name of the test
+	Type     TestType          // type of the test
+	Subtest  bool
 }
 
 func (t *gobTest) result() Result {

gopls/internal/golang/code_lens.go

@@ -19,6 +19,7 @@ import (
 	"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/cache/testfuncs"
 	"golang.org/x/tools/gopls/internal/file"
 	"golang.org/x/tools/gopls/internal/protocol"
 	"golang.org/x/tools/gopls/internal/protocol/command"
@@ -213,6 +214,29 @@ func regenerateCgoLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Ha
 }
 
 func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle) ([]protocol.CodeLens, error) {
+	matches, err := matchFunctionsWithTests(ctx, snapshot, fh)
+	if err != nil {
+		return nil, err
+	}
+
+	lenses := make([]protocol.CodeLens, 0, len(matches))
+	for _, t := range matches {
+		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
+}
+
+type TestMatch struct {
+	FuncPos protocol.Position  // function position
+	Name    string             // test name
+	Loc     protocol.Location  // test location
+	Type    testfuncs.TestType // test type
+}
+
+func matchFunctionsWithTests(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle) (matches []TestMatch, err error) {
 	if strings.HasSuffix(fh.URI().Path(), "_test.go") {
 		// Ignore test files.
 		return nil, nil
@@ -238,7 +262,12 @@ func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	if err != nil {
 		return nil, fmt.Errorf("couldn't parse file: %w", err)
 	}
-	funcPos := make(map[string]protocol.Position)
+
+	type Func struct {
+		Name string
+		Pos  protocol.Position
+	}
+	var fileFuncs []Func
 	for _, d := range pgf.File.Decls {
 		fn, ok := d.(*ast.FuncDecl)
 		if !ok {
@@ -254,32 +283,11 @@ func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 			_, 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
+		fileFuncs = append(fileFuncs, Func{
+			Name: name,
+			Pos:  rng.Start,
+		})
 	}
-	var matchedTests []Test
 
 	pkgIDs := make([]PackageID, 0, len(testPackages))
 	for _, pkg := range testPackages {
@@ -291,48 +299,54 @@ func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	}
 	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 test.Subtest {
+				continue
 			}
-			if testType == 0 {
-				continue // unknown type
+			potentialFuncNames := getPotentialFuncNames(test)
+			if len(potentialFuncNames) == 0 {
+				continue
 			}
-			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)]
+
+			var matchedFunc Func
+			for _, fn := range fileFuncs {
+				var matched bool
+				for _, n := range potentialFuncNames {
+					// Check the prefix to be able to match 'TestDeletePanics' with 'Delete'.
+					if strings.HasPrefix(n, fn.Name) {
+						matched = true
+						break
+					}
+				}
+				if !matched {
+					continue
+				}
+
+				// Use the most specific function:
+				//
+				//   - match 'TestDelete', 'TestDeletePanics' with 'Delete'
+				//   - match 'TestDeleteFunc', 'TestDeleteFuncClearTail' with 'DeleteFunc', not 'Delete'
+				if len(matchedFunc.Name) < len(fn.Name) {
+					matchedFunc = fn
+				}
 			}
-			if ok {
+			if matchedFunc.Name != "" {
 				loc := test.Location
 				loc.Range.End = loc.Range.Start // move cursor to the test's beginning
 
-				matchedTests = append(matchedTests, Test{
-					FuncPos: pos,
+				matches = append(matches, TestMatch{
+					FuncPos: matchedFunc.Pos,
 					Name:    test.Name,
 					Loc:     loc,
-					Type:    testType,
+					Type:    test.Type,
 				})
 			}
 		}
 	}
-	if len(matchedTests) == 0 {
+	if len(matches) == 0 {
 		return nil, nil
 	}
 
-	slices.SortFunc(matchedTests, func(a, b Test) int {
+	slices.SortFunc(matches, func(a, b TestMatch) int {
 		if v := protocol.ComparePosition(a.FuncPos, b.FuncPos); v != 0 {
 			return v
 		}
@@ -341,13 +355,31 @@ func goToTestCodeLens(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 		}
 		return cmp.Compare(a.Name, b.Name)
 	})
+	return matches, nil
+}
 
-	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),
-		})
+func getPotentialFuncNames(test testfuncs.Result) []string {
+	var name string
+	switch test.Type {
+	case testfuncs.TypeTest:
+		name = strings.TrimPrefix(test.Name, "Test")
+	case testfuncs.TypeBenchmark:
+		name = strings.TrimPrefix(test.Name, "Benchmark")
+	case testfuncs.TypeFuzz:
+		name = strings.TrimPrefix(test.Name, "Fuzz")
+	case testfuncs.TypeExample:
+		name = strings.TrimPrefix(test.Name, "Example")
+	}
+	if name == "" {
+		return nil
+	}
+	name = strings.TrimPrefix(name, "_")
+
+	lowerCasedName := []rune(name)
+	lowerCasedName[0] = unicode.ToLower(lowerCasedName[0])
+
+	return []string{
+		name,                   // 'Foo' for 'TestFoo', 'foo' for 'Test_foo'
+		string(lowerCasedName), // 'foo' for 'TestFoo'
 	}
-	return lenses, nil
 }

gopls/internal/golang/codeaction.go

@@ -266,6 +266,7 @@ var codeActionProducers = [...]codeActionProducer{
 	{kind: settings.RefactorRewriteEliminateDotImport, fn: refactorRewriteEliminateDotImport, needPkg: true},
 	{kind: settings.RefactorRewriteAddTags, fn: refactorRewriteAddStructTags, needPkg: true},
 	{kind: settings.RefactorRewriteRemoveTags, fn: refactorRewriteRemoveStructTags, needPkg: true},
+	{kind: settings.GoToTest, fn: goToTest, needPkg: true},
 	{kind: settings.GoplsDocFeatures, fn: goplsDocFeatures}, // offer this one last (#72742)
 
 	// Note: don't forget to update the allow-list in Server.CodeAction
@@ -1129,3 +1130,34 @@ func toggleCompilerOptDetails(ctx context.Context, req *codeActionsRequest) erro
 	}
 	return nil
 }
+
+// goToTest produces "Go to TestXxx" code action.
+// See [server.commandHandler.GoToTest] for command implementation.
+func goToTest(ctx context.Context, req *codeActionsRequest) error {
+	// TODO: add tests.
+
+	path, _ := astutil.PathEnclosingInterval(req.pgf.File, req.start, req.end)
+	if len(path) < 2 {
+		return nil
+	}
+	fn, ok := path[len(path)-2].(*ast.FuncDecl)
+	if !ok {
+		return nil
+	}
+	fnRng, err := req.pgf.NodeRange(fn)
+	if err != nil {
+		return fmt.Errorf("couldn't get node range: %w", err)
+	}
+
+	matches, err := matchFunctionsWithTests(ctx, req.snapshot, req.fh)
+	if err != nil {
+		return err
+	}
+	for _, m := range matches {
+		if m.FuncPos == fnRng.Start {
+			cmd := command.NewGoToTestCommand("Go to "+m.Name, m.Loc)
+			req.addCommandAction(cmd, false)
+		}
+	}
+	return nil
+}

gopls/internal/settings/codeactionkind.go

@@ -83,6 +83,7 @@ const (
 	GoToggleCompilerOptDetails protocol.CodeActionKind = "source.toggleCompilerOptDetails"
 	AddTest                    protocol.CodeActionKind = "source.addTest"
 	OrganizeImports            protocol.CodeActionKind = "source.organizeImports"
+	GoToTest                   protocol.CodeActionKind = "source.go_to_test"
 
 	// gopls
 	GoplsDocFeatures protocol.CodeActionKind = "gopls.doc.features"