gopls/internal/lsp: add selection range request

Fixes golang/go#36679

gopls/internal/lsp/cmd/test/cmdtest.go

@@ -118,6 +118,10 @@ func (r *runner) InlayHints(t *testing.T, spn span.Span) {
 	// TODO: inlayHints not supported on command line
 }
 
+func (r *runner) SelectionRanges(t *testing.T, spn span.Span, exp span.Span) {
+	// TODO: selectionRanges not supported on command line
+}
+
 func (r *runner) runGoplsCmd(t testing.TB, args ...string) (string, string) {
 	rStdout, wStdout, err := os.Pipe()
 	if err != nil {

gopls/internal/lsp/lsp_test.go

@@ -1282,6 +1282,34 @@ func (r *runner) AddImport(t *testing.T, uri span.URI, expectedImport string) {
 	}
 }
 
+func (r *runner) SelectionRanges(t *testing.T, spn span.Span, exp span.Span) {
+	sm, err := r.data.Mapper(spn.URI())
+	if err != nil {
+		t.Fatal(err)
+	}
+	loc, err := sm.Location(spn)
+
+	ranges, err := r.server.selectionRange(r.ctx, &protocol.SelectionRangeParams{
+		TextDocument: protocol.TextDocumentIdentifier{
+			URI: protocol.URIFromSpanURI(spn.URI()),
+		},
+		Positions: []protocol.Position{loc.Range.Start, loc.Range.End},
+	})
+	if err != nil {
+		t.Error(err)
+	}
+
+	if len(ranges) != 1 {
+		t.Error(ranges)
+	}
+
+	exploc, err := sm.Location(exp)
+
+	if ranges[0].Range != exploc.Range {
+		t.Errorf("expected %v, actual %v", exploc.Range, ranges[0].Range)
+	}
+}
+
 func TestBytesOffset(t *testing.T) {
 	tests := []struct {
 		text string

gopls/internal/lsp/selection_range.go

@@ -0,0 +1,55 @@
+package lsp
+
+import (
+	"context"
+	"fmt"
+
+	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/gopls/internal/lsp/protocol"
+	"golang.org/x/tools/gopls/internal/lsp/source"
+	"golang.org/x/tools/internal/event"
+)
+
+func (s *Server) selectionRange(ctx context.Context, params *protocol.SelectionRangeParams) ([]protocol.SelectionRange, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.documentSymbol")
+	defer done()
+
+	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
+	defer release()
+	if !ok {
+		return nil, err
+	}
+
+	pgf, err := snapshot.ParseGo(ctx, fh, source.ParseFull)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(params.Positions) != 2 {
+		return nil, fmt.Errorf("expected 2 positions, received %d", len(params.Positions))
+	}
+
+	start, err := pgf.Mapper.Pos(params.Positions[0])
+	if err != nil {
+		return nil, err
+	}
+
+	end, err := pgf.Mapper.Pos(params.Positions[1])
+	if err != nil {
+		return nil, err
+	}
+
+	path, _ := astutil.PathEnclosingInterval(pgf.File, start, end)
+
+	n := path[0]
+	if len(path) >= 2 && n.Pos() == start && n.End() == end {
+		n = path[1]
+	}
+
+	newSelection, err := pgf.Mapper.PosRange(n.Pos(), n.End())
+	if err != nil {
+		return nil, err
+	}
+
+	return []protocol.SelectionRange{{Range: newSelection}}, nil
+}

gopls/internal/lsp/server_gen.go

@@ -244,8 +244,8 @@ func (s *Server) ResolveWorkspaceSymbol(context.Context, *protocol.WorkspaceSymb
 	return nil, notImplemented("ResolveWorkspaceSymbol")
 }
 
-func (s *Server) SelectionRange(context.Context, *protocol.SelectionRangeParams) ([]protocol.SelectionRange, error) {
-	return nil, notImplemented("SelectionRange")
+func (s *Server) SelectionRange(ctx context.Context, params *protocol.SelectionRangeParams) ([]protocol.SelectionRange, error) {
+	return s.selectionRange(ctx, params)
 }
 
 func (s *Server) SemanticTokensFull(ctx context.Context, p *protocol.SemanticTokensParams) (*protocol.SemanticTokens, error) {

gopls/internal/lsp/source/source_test.go

@@ -497,6 +497,10 @@ func (r *runner) SemanticTokens(t *testing.T, spn span.Span) {
 	t.Skip("nothing to test in source")
 }
 
+func (r *runner) SelectionRanges(t *testing.T, spn, exp span.Span) {
+	t.Skip("nothing to test in source")
+}
+
 func (r *runner) Import(t *testing.T, spn span.Span) {
 	fh, err := r.snapshot.GetFile(r.ctx, spn.URI())
 	if err != nil {

gopls/internal/lsp/testdata/selectionrange/foo.go

@@ -0,0 +1,13 @@
+package foo
+
+import "time"
+
+func Bar(x, y int, t time.Time) int {
+	zs := []int{1, 2, 3} //@selectionrange("1, 2", "[]int{1, 2, 3}")
+
+	for _, z := range zs {
+		x = x + z + y + zs[1]
+	}
+
+	return x + y //@selectionrange("x + ", "x + y")
+}

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

@@ -28,4 +28,5 @@ WorkspaceSymbolsCount = 20
 SignaturesCount = 33
 LinksCount = 7
 ImplementationsCount = 14
+SelectionRangesCount = 2
 

gopls/internal/lsp/testdata/summary_go1.18.txt.golden

@@ -28,4 +28,5 @@ WorkspaceSymbolsCount = 20
 SignaturesCount = 33
 LinksCount = 7
 ImplementationsCount = 14
+SelectionRangesCount = 2
 

gopls/internal/lsp/tests/tests.go

@@ -93,6 +93,7 @@ type Signatures = map[span.Span]*protocol.SignatureHelp
 type Links = map[span.URI][]Link
 type AddImport = map[span.URI]string
 type Hovers = map[span.Span]string
+type SelectionRanges = map[span.Span]span.Span
 
 type Data struct {
 	Config                   packages.Config
@@ -128,6 +129,7 @@ type Data struct {
 	Links                    Links
 	AddImport                AddImport
 	Hovers                   Hovers
+	SelectionRanges          SelectionRanges
 
 	fragments map[string]string
 	dir       string
@@ -177,6 +179,7 @@ type Tests interface {
 	Link(*testing.T, span.URI, []Link)
 	AddImport(*testing.T, span.URI, string)
 	Hover(*testing.T, span.Span, string)
+	SelectionRanges(*testing.T, span.Span, span.Span)
 }
 
 type Definition struct {
@@ -334,6 +337,7 @@ func load(t testing.TB, mode string, dir string) *Data {
 		Links:                    make(Links),
 		AddImport:                make(AddImport),
 		Hovers:                   make(Hovers),
+		SelectionRanges:          make(SelectionRanges),
 
 		dir:       dir,
 		fragments: map[string]string{},
@@ -499,6 +503,7 @@ func load(t testing.TB, mode string, dir string) *Data {
 		"incomingcalls":   datum.collectIncomingCalls,
 		"outgoingcalls":   datum.collectOutgoingCalls,
 		"addimport":       datum.collectAddImports,
+		"selectionrange":  datum.collectSelectionRanges,
 	}); err != nil {
 		t.Fatal(err)
 	}
@@ -947,6 +952,15 @@ func Run(t *testing.T, tests Tests, data *Data) {
 		}
 	})
 
+	t.Run("SelectionRanges", func(t *testing.T) {
+		t.Helper()
+		for span, exp := range data.SelectionRanges {
+			t.Run(SpanName(span), func(t *testing.T) {
+				tests.SelectionRanges(t, span, exp)
+			})
+		}
+	})
+
 	if *UpdateGolden {
 		for _, golden := range data.golden {
 			if !golden.Modified {
@@ -1039,6 +1053,7 @@ func checkData(t *testing.T, data *Data) {
 	fmt.Fprintf(buf, "SignaturesCount = %v\n", len(data.Signatures))
 	fmt.Fprintf(buf, "LinksCount = %v\n", linksCount)
 	fmt.Fprintf(buf, "ImplementationsCount = %v\n", len(data.Implementations))
+	fmt.Fprintf(buf, "SelectionRangesCount = %v\n", len(data.SelectionRanges))
 
 	want := string(data.Golden(t, "summary", summaryFile, func() ([]byte, error) {
 		return buf.Bytes(), nil
@@ -1240,6 +1255,12 @@ func (data *Data) collectDefinitions(src, target span.Span) {
 	}
 }
 
+func (data *Data) collectSelectionRanges(src, exp span.Span) {
+	if _, ok := data.SelectionRanges[src]; !ok {
+		data.SelectionRanges[src] = exp
+	}
+}
+
 func (data *Data) collectImplementations(src span.Span, targets []span.Span) {
 	data.Implementations[src] = targets
 }