internal/lsp: add support for formatting go.work files

Wired through support for calling x/mod's go.work formatter on go.work
files into LSP. Tested it by hand in editor using the "Format Document"
command. Added a test case to workspace_test regtest, though I'm not
totally sure the test is correct.

For golang/go#50930

Change-Id: Ied052ded514bb36f561737698f0e2d7b488158e7
Reviewed-on: https://go-review.googlesource.com/c/tools/+/383774
Trust: Michael Matloob <matloob@golang.org>
Run-TryBot: Michael Matloob <matloob@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>

gopls/internal/regtest/workspace/workspace_test.go

@@ -771,6 +771,28 @@ use (
 		if err := checkHelloLocation("b.com@v1.2.3/b/b.go"); err != nil {
 			t.Fatal(err)
 		}
+
+		// Test Formatting.
+		env.SetBufferContent("go.work", `go 1.18
+  use      (
+
+
+
+		./moda/a
+)
+`) // TODO(matloob): For some reason there's a "start position 7:0 is out of bounds" error when the ")" is on the last character/line in the file. Rob probably knows what's going on.
+		env.SaveBuffer("go.work")
+		env.Await(env.DoneWithSave())
+		gotWorkContents := env.ReadWorkspaceFile("go.work")
+		wantWorkContents := `go 1.18
+
+use (
+	./moda/a
+)
+`
+		if gotWorkContents != wantWorkContents {
+			t.Fatalf("formatted contents of workspace: got %q; want %q", gotWorkContents, wantWorkContents)
+		}
 	})
 }
 

internal/lsp/cache/mod.go

@@ -101,6 +101,84 @@ func (s *snapshot) ParseMod(ctx context.Context, modFH source.FileHandle) (*sour
 	return pmh.parse(ctx, s)
 }
 
+type parseWorkHandle struct {
+	handle *memoize.Handle
+}
+
+type parseWorkData struct {
+	parsed *source.ParsedWorkFile
+
+	// err is any error encountered while parsing the file.
+	err error
+}
+
+func (mh *parseWorkHandle) parse(ctx context.Context, snapshot *snapshot) (*source.ParsedWorkFile, error) {
+	v, err := mh.handle.Get(ctx, snapshot.generation, snapshot)
+	if err != nil {
+		return nil, err
+	}
+	data := v.(*parseWorkData)
+	return data.parsed, data.err
+}
+
+func (s *snapshot) ParseWork(ctx context.Context, modFH source.FileHandle) (*source.ParsedWorkFile, error) {
+	if handle := s.getParseWorkHandle(modFH.URI()); handle != nil {
+		return handle.parse(ctx, s)
+	}
+	h := s.generation.Bind(modFH.FileIdentity(), func(ctx context.Context, _ memoize.Arg) interface{} {
+		_, done := event.Start(ctx, "cache.ParseModHandle", tag.URI.Of(modFH.URI()))
+		defer done()
+
+		contents, err := modFH.Read()
+		if err != nil {
+			return &parseModData{err: err}
+		}
+		m := &protocol.ColumnMapper{
+			URI:       modFH.URI(),
+			Converter: span.NewContentConverter(modFH.URI().Filename(), contents),
+			Content:   contents,
+		}
+		file, parseErr := modfile.ParseWork(modFH.URI().Filename(), contents, nil)
+		// Attempt to convert the error to a standardized parse error.
+		var parseErrors []*source.Diagnostic
+		if parseErr != nil {
+			mfErrList, ok := parseErr.(modfile.ErrorList)
+			if !ok {
+				return &parseModData{err: fmt.Errorf("unexpected parse error type %v", parseErr)}
+			}
+			for _, mfErr := range mfErrList {
+				rng, err := rangeFromPositions(m, mfErr.Pos, mfErr.Pos)
+				if err != nil {
+					return &parseModData{err: err}
+				}
+				parseErrors = []*source.Diagnostic{{
+					URI:      modFH.URI(),
+					Range:    rng,
+					Severity: protocol.SeverityError,
+					Source:   source.ParseError,
+					Message:  mfErr.Err.Error(),
+				}}
+			}
+		}
+		return &parseWorkData{
+			parsed: &source.ParsedWorkFile{
+				URI:         modFH.URI(),
+				Mapper:      m,
+				File:        file,
+				ParseErrors: parseErrors,
+			},
+			err: parseErr,
+		}
+	}, nil)
+
+	pwh := &parseWorkHandle{handle: h}
+	s.mu.Lock()
+	s.parseWorkHandles[modFH.URI()] = pwh
+	s.mu.Unlock()
+
+	return pwh.parse(ctx, s)
+}
+
 // goSum reads the go.sum file for the go.mod file at modURI, if it exists. If
 // it doesn't exist, it returns nil.
 func (s *snapshot) goSum(ctx context.Context, modURI span.URI) []byte {

internal/lsp/cache/session.go

@@ -243,6 +243,7 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 		workspacePackages: make(map[PackageID]PackagePath),
 		unloadableFiles:   make(map[span.URI]struct{}),
 		parseModHandles:   make(map[span.URI]*parseModHandle),
+		parseWorkHandles:  make(map[span.URI]*parseWorkHandle),
 		modTidyHandles:    make(map[span.URI]*modTidyHandle),
 		modWhyHandles:     make(map[span.URI]*modWhyHandle),
 		workspace:         workspace,

internal/lsp/cache/snapshot.go

@@ -102,10 +102,14 @@ type snapshot struct {
 	// unloadableFiles keeps track of files that we've failed to load.
 	unloadableFiles map[span.URI]struct{}
 
-	// parseModHandles keeps track of any ParseModHandles for the snapshot.
+	// parseModHandles keeps track of any parseModHandles for the snapshot.
 	// The handles need not refer to only the view's go.mod file.
 	parseModHandles map[span.URI]*parseModHandle
 
+	// parseWorkHandles keeps track of any parseWorkHandles for the snapshot.
+	// The handles need not refer to only the view's go.work file.
+	parseWorkHandles map[span.URI]*parseWorkHandle
+
 	// Preserve go.mod-related handles to avoid garbage-collecting the results
 	// of various calls to the go command. The handles need not refer to only
 	// the view's go.mod file.
@@ -688,6 +692,12 @@ func (s *snapshot) getParseModHandle(uri span.URI) *parseModHandle {
 	return s.parseModHandles[uri]
 }
 
+func (s *snapshot) getParseWorkHandle(uri span.URI) *parseWorkHandle {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.parseWorkHandles[uri]
+}
+
 func (s *snapshot) getModWhyHandle(uri span.URI) *modWhyHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -1729,6 +1739,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		workspacePackages: make(map[PackageID]PackagePath, len(s.workspacePackages)),
 		unloadableFiles:   make(map[span.URI]struct{}, len(s.unloadableFiles)),
 		parseModHandles:   make(map[span.URI]*parseModHandle, len(s.parseModHandles)),
+		parseWorkHandles:  make(map[span.URI]*parseWorkHandle, len(s.parseWorkHandles)),
 		modTidyHandles:    make(map[span.URI]*modTidyHandle, len(s.modTidyHandles)),
 		modWhyHandles:     make(map[span.URI]*modWhyHandle, len(s.modWhyHandles)),
 		knownSubdirs:      make(map[span.URI]struct{}, len(s.knownSubdirs)),
@@ -1763,6 +1774,10 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 	for k, v := range s.parseModHandles {
 		result.parseModHandles[k] = v
 	}
+	// Copy all of the parseWorkHandles.
+	for k, v := range s.parseWorkHandles {
+		result.parseWorkHandles[k] = v
+	}
 
 	for k, v := range s.goFiles {
 		if _, ok := changes[k.file.URI]; ok {
@@ -1853,9 +1868,8 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 				delete(result.modWhyHandles, k)
 			}
 		}
-		if isGoMod(uri) {
-			delete(result.parseModHandles, uri)
-		}
+		delete(result.parseModHandles, uri)
+		delete(result.parseWorkHandles, uri)
 		// Handle the invalidated file; it may have new contents or not exist.
 		if !change.exists {
 			delete(result.files, uri)
@@ -2060,6 +2074,9 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 	for _, v := range result.parseModHandles {
 		newGen.Inherit(v.handle)
 	}
+	for _, v := range result.parseWorkHandles {
+		newGen.Inherit(v.handle)
+	}
 	// Don't bother copying the importedBy graph,
 	// as it changes each time we update metadata.
 

internal/lsp/cache/workspace.go

@@ -439,7 +439,7 @@ func isGoMod(uri span.URI) bool {
 }
 
 func isGoSum(uri span.URI) bool {
-	return filepath.Base(uri.Filename()) == "go.sum"
+	return filepath.Base(uri.Filename()) == "go.sum" || filepath.Base(uri.Filename()) == "go.work.sum"
 }
 
 // fileExists reports if the file uri exists within source.

internal/lsp/fake/editor.go

@@ -434,10 +434,11 @@ func (e *Editor) createBuffer(ctx context.Context, path string, dirty bool, cont
 }
 
 var defaultFileAssociations = map[string]*regexp.Regexp{
-	"go":     regexp.MustCompile(`^.*\.go$`), // '$' is important: don't match .gotmpl!
-	"go.mod": regexp.MustCompile(`^go\.mod$`),
-	"go.sum": regexp.MustCompile(`^go\.sum$`),
-	"gotmpl": regexp.MustCompile(`^.*tmpl$`),
+	"go":      regexp.MustCompile(`^.*\.go$`), // '$' is important: don't match .gotmpl!
+	"go.mod":  regexp.MustCompile(`^go\.mod$`),
+	"go.sum":  regexp.MustCompile(`^go(\.work)?\.sum$`),
+	"go.work": regexp.MustCompile(`^go\.work$`),
+	"gotmpl":  regexp.MustCompile(`^.*tmpl$`),
 }
 
 func (e *Editor) languageID(p string) string {

internal/lsp/format.go

@@ -10,6 +10,7 @@ import (
 	"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/work"
 )
 
 func (s *Server) formatting(ctx context.Context, params *protocol.DocumentFormattingParams) ([]protocol.TextEdit, error) {
@@ -23,6 +24,8 @@ func (s *Server) formatting(ctx context.Context, params *protocol.DocumentFormat
 		return mod.Format(ctx, snapshot, fh)
 	case source.Go:
 		return source.Format(ctx, snapshot, fh)
+	case source.Work:
+		return work.Format(ctx, snapshot, fh)
 	}
 	return nil, nil
 }

internal/lsp/source/options.go

@@ -103,6 +103,7 @@ func DefaultOptions() *Options {
 						protocol.SourceOrganizeImports: true,
 						protocol.QuickFix:              true,
 					},
+					Work: {},
 					Sum:  {},
 					Tmpl: {},
 				},

internal/lsp/source/view.go

@@ -129,6 +129,9 @@ type Snapshot interface {
 	// GoModForFile returns the URI of the go.mod file for the given URI.
 	GoModForFile(uri span.URI) span.URI
 
+	// ParseWork is used to parse go.work files.
+	ParseWork(ctx context.Context, fh FileHandle) (*ParsedWorkFile, error)
+
 	// BuiltinFile returns information about the special builtin package.
 	BuiltinFile(ctx context.Context) (*ParsedGoFile, error)
 
@@ -293,6 +296,14 @@ type ParsedModule struct {
 	ParseErrors []*Diagnostic
 }
 
+// A ParsedWorkFile contains the results of parsing a go.work file.
+type ParsedWorkFile struct {
+	URI         span.URI
+	File        *modfile.WorkFile
+	Mapper      *protocol.ColumnMapper
+	ParseErrors []*Diagnostic
+}
+
 // A TidiedModule contains the results of running `go mod tidy` on a module.
 type TidiedModule struct {
 	// Diagnostics representing changes made by `go mod tidy`.

internal/lsp/tests/tests.go

@@ -252,6 +252,7 @@ func DefaultOptions(o *source.Options) {
 			protocol.SourceOrganizeImports: true,
 		},
 		source.Sum:  {},
+		source.Work: {},
 		source.Tmpl: {},
 	}
 	o.UserOptions.Codelenses[string(command.Test)] = true

internal/lsp/work/format.go

@@ -0,0 +1,31 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package work
+
+import (
+	"context"
+
+	"golang.org/x/mod/modfile"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/lsp/protocol"
+	"golang.org/x/tools/internal/lsp/source"
+)
+
+func Format(ctx context.Context, snapshot source.Snapshot, fh source.FileHandle) ([]protocol.TextEdit, error) {
+	ctx, done := event.Start(ctx, "work.Format")
+	defer done()
+
+	pw, err := snapshot.ParseWork(ctx, fh)
+	if err != nil {
+		return nil, err
+	}
+	formatted := modfile.Format(pw.File.Syntax)
+	// Calculate the edits to be made due to the change.
+	diff, err := snapshot.View().Options().ComputeEdits(fh.URI(), string(pw.Mapper.Content), string(formatted))
+	if err != nil {
+		return nil, err
+	}
+	return source.ToProtocolEdits(pw.Mapper, diff)
+}