gopls/internal/golang: definition support return statements

Enables you to use the definition operation on a return token
and brings you to the return parameters.

Updates golang/go#70462

Change-Id: I080ff37aad5ac1498e320429fad6dd4b043f02c8
Reviewed-on: https://go-review.googlesource.com/c/tools/+/636855
Auto-Submit: Madeline Kalil <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>

Author: madelinekalil

gopls/doc/features/navigation.md

@@ -24,6 +24,7 @@ A definition query also works in these unexpected places:
   it returns the location of the embedded file.
 - On the declaration of a non-Go function (a `func` with no body),
   it returns the location of the assembly implementation, if any,
+- On a **return statement**, it returns the location of the function's result variables.
 
 <!-- On a built-in symbol such as `append` or `unsafe.Pointer`, `definition` reports
 the location of the declaration in the builtin or unsafe pseudo-packages,

gopls/doc/release/v0.18.0.md

@@ -28,3 +28,10 @@ func (C[T]) Pop() (T, bool) { ... }
 var _ Stack[int] = C[int]{}
 ```
 
+## Improvements to "Definition"
+
+The Definition query now supports additional locations:
+
+- When invoked on a return statement, it reports the location
+  of the function's result variables.
+

gopls/internal/golang/definition.go

@@ -15,12 +15,13 @@ import (
 	"regexp"
 	"strings"
 
+	"golang.org/x/tools/go/ast/astutil"
 	"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/util/astutil"
+	goplsastutil "golang.org/x/tools/gopls/internal/util/astutil"
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/internal/event"
 )
@@ -84,6 +85,32 @@ func Definition(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, p
 		return locations, err // may be success or failure
 	}
 
+	// Handle the case where the cursor is on a return statement by jumping to the result variables.
+	path, _ := astutil.PathEnclosingInterval(pgf.File, pos, pos)
+	if is[*ast.ReturnStmt](path[0]) {
+		var funcType *ast.FuncType
+		for _, n := range path {
+			switch n := n.(type) {
+			case *ast.FuncLit:
+				funcType = n.Type
+			case *ast.FuncDecl:
+				funcType = n.Type
+			}
+			if funcType != nil {
+				break
+			}
+		}
+		// Inv: funcType != nil, as a return stmt cannot appear outside a function.
+		if funcType.Results == nil {
+			return nil, nil // no result variables
+		}
+		loc, err := pgf.NodeLocation(funcType.Results)
+		if err != nil {
+			return nil, err
+		}
+		return []protocol.Location{loc}, nil
+	}
+
 	// The general case: the cursor is on an identifier.
 	_, obj, _ := referencedObject(pkg, pgf, pos)
 	if obj == nil {
@@ -102,7 +129,7 @@ func Definition(ctx context.Context, snapshot *cache.Snapshot, fh file.Handle, p
 	for _, decl := range pgf.File.Decls {
 		if decl, ok := decl.(*ast.FuncDecl); ok &&
 			decl.Body == nil &&
-			astutil.NodeContains(decl.Name, pos) {
+			goplsastutil.NodeContains(decl.Name, pos) {
 			return nonGoDefinition(ctx, snapshot, pkg, decl.Name.Name)
 		}
 	}

gopls/internal/test/marker/testdata/definition/return.txt

@@ -0,0 +1,23 @@
+This test checks definition operations in function return statements. 
+Go to definition on 'return' should go to the result parameter list.
+
+-- go.mod --
+module mod.com
+
+go 1.18
+
+-- a/a.go --
+package a
+
+func Hi() string { //@loc(HiReturn, "string")
+    return "Hello" //@def("return", HiReturn)
+}
+
+func Bye() (int, int, int) { //@loc(ByeReturn, "(int, int, int)")
+    return 1, 2, 3 //@def("return", ByeReturn)
+}
+
+func TestLit() {
+    f := func(a, b int) bool { return a*b < 100 } //@loc(FuncLitReturn, "bool"),def("return", FuncLitReturn)
+    f(1, 2)
+}