internal/lsp: start parsing go.work file like gopls.mod file

Allow users to experiment with golang/go#45713 by adding experimental
support for the go.work file. We handle it like a special case, very
similar to the current gopls.mod file mechanism. The behavior is
undefined if both a gopls.mod and go.work file exist. Ultimately, we
will deprecate support for the gopls.mod file if the go.work file
proposal is accepted, so I don't think it's important to be careful
about handling both simultaneously.

Change-Id: Id822aeec953fc1d4acca343b742afea899dc70ff
Reviewed-on: https://go-review.googlesource.com/c/tools/+/328334
Trust: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
diff --git a/gopls/internal/regtest/workspace/workspace_test.go b/gopls/internal/regtest/workspace/workspace_test.go
index 6a7b657..4c5d1fc 100644
--- a/gopls/internal/regtest/workspace/workspace_test.go
+++ b/gopls/internal/regtest/workspace/workspace_test.go
@@ -611,6 +611,132 @@
 	})
 }
 
+func TestUseGoWork(t *testing.T) {
+	// This test validates certain functionality related to using a go.work
+	// file to specify workspace modules.
+	testenv.NeedsGo1Point(t, 14)
+	const multiModule = `
+-- moda/a/go.mod --
+module a.com
+
+require b.com v1.2.3
+-- moda/a/go.sum --
+b.com v1.2.3 h1:tXrlXP0rnjRpKNmkbLYoWBdq0ikb3C3bKK9//moAWBI=
+b.com v1.2.3/go.mod h1:D+J7pfFBZK5vdIdZEFquR586vKKIkqG7Qjw9AxG5BQ8=
+-- moda/a/a.go --
+package a
+
+import (
+	"b.com/b"
+)
+
+func main() {
+	var x int
+	_ = b.Hello()
+}
+-- modb/go.mod --
+module b.com
+
+require example.com v1.2.3
+-- modb/go.sum --
+example.com v1.2.3 h1:Yryq11hF02fEf2JlOS2eph+ICE2/ceevGV3C9dl5V/c=
+example.com v1.2.3/go.mod h1:Y2Rc5rVWjWur0h3pd9aEvK5Pof8YKDANh9gHA2Maujo=
+-- modb/b/b.go --
+package b
+
+func Hello() int {
+	var x int
+}
+-- go.work --
+go 1.17
+
+directory (
+	./moda/a
+)
+`
+	WithOptions(
+		ProxyFiles(workspaceModuleProxy),
+		Modes(Experimental),
+	).Run(t, multiModule, func(t *testing.T, env *Env) {
+		// Initially, the gopls.mod should cause only the a.com module to be
+		// loaded. Validate this by jumping to a definition in b.com and ensuring
+		// that we go to the module cache.
+		env.OpenFile("moda/a/a.go")
+		env.Await(env.DoneWithOpen())
+
+		// To verify which modules are loaded, we'll jump to the definition of
+		// b.Hello.
+		checkHelloLocation := func(want string) error {
+			location, _ := env.GoToDefinition("moda/a/a.go", env.RegexpSearch("moda/a/a.go", "Hello"))
+			if !strings.HasSuffix(location, want) {
+				return fmt.Errorf("expected %s, got %v", want, location)
+			}
+			return nil
+		}
+
+		// Initially this should be in the module cache, as b.com is not replaced.
+		if err := checkHelloLocation("b.com@v1.2.3/b/b.go"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Now, modify the gopls.mod file on disk to activate the b.com module in
+		// the workspace.
+		env.WriteWorkspaceFile("go.work", `
+go 1.17
+
+directory (
+	./moda/a
+	./modb
+)
+`)
+		env.Await(env.DoneWithChangeWatchedFiles())
+		// Check that go.mod diagnostics picked up the newly active mod file.
+		// The local version of modb has an extra dependency we need to download.
+		env.OpenFile("modb/go.mod")
+		env.Await(env.DoneWithOpen())
+
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexpWithMessage("modb/go.mod", `require example.com v1.2.3`, "has not been downloaded"),
+				ReadDiagnostics("modb/go.mod", &d),
+			),
+		)
+		env.ApplyQuickFixes("modb/go.mod", d.Diagnostics)
+		env.Await(env.DiagnosticAtRegexp("modb/b/b.go", "x"))
+		// Jumping to definition should now go to b.com in the workspace.
+		if err := checkHelloLocation("modb/b/b.go"); err != nil {
+			t.Fatal(err)
+		}
+
+		// Now, let's modify the gopls.mod *overlay* (not on disk), and verify that
+		// this change is only picked up once it is saved.
+		env.OpenFile("go.work")
+		env.Await(env.DoneWithOpen())
+		env.SetBufferContent("go.work", `go 1.17
+
+directory (
+	./moda/a
+)`)
+
+		// Editing the gopls.mod removes modb from the workspace modules, and so
+		// should clear outstanding diagnostics...
+		env.Await(OnceMet(
+			env.DoneWithChange(),
+			EmptyDiagnostics("modb/go.mod"),
+		))
+		// ...but does not yet cause a workspace reload, so we should still jump to modb.
+		if err := checkHelloLocation("modb/b/b.go"); err != nil {
+			t.Fatal(err)
+		}
+		// Saving should reload the workspace.
+		env.SaveBufferWithoutActions("go.work")
+		if err := checkHelloLocation("b.com@v1.2.3/b/b.go"); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
 func TestNonWorkspaceFileCreation(t *testing.T) {
 	testenv.NeedsGo1Point(t, 13)
 
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index 7b2c882..207be3a 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -1528,9 +1528,25 @@
 	}
 }
 
+// unappliedChanges is a file source that handles an uncloned snapshot.
+type unappliedChanges struct {
+	originalSnapshot *snapshot
+	changes          map[span.URI]*fileChange
+}
+
+func (ac *unappliedChanges) GetFile(ctx context.Context, uri span.URI) (source.FileHandle, error) {
+	if c, ok := ac.changes[uri]; ok {
+		return c.fileHandle, nil
+	}
+	return ac.originalSnapshot.GetFile(ctx, uri)
+}
+
 func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileChange, forceReloadMetadata bool) (*snapshot, bool) {
 	var vendorChanged bool
-	newWorkspace, workspaceChanged, workspaceReload := s.workspace.invalidate(ctx, changes)
+	newWorkspace, workspaceChanged, workspaceReload := s.workspace.invalidate(ctx, changes, &unappliedChanges{
+		originalSnapshot: s,
+		changes:          changes,
+	})
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index fea4e81..f6aefd8 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -393,12 +393,16 @@
 	if v.knownFile(c.URI) {
 		return true
 	}
-	// The gopls.mod may not be "known" because we first access it through the
-	// session. As a result, treat changes to the view's gopls.mod file as
-	// always relevant, even if they are only on-disk changes.
-	// TODO(rstambler): Make sure the gopls.mod is always known to the view.
-	if c.URI == goplsModURI(v.rootURI) {
-		return true
+	// The go.work/gopls.mod may not be "known" because we first access it
+	// through the session. As a result, treat changes to the view's go.work or
+	// gopls.mod file as always relevant, even if they are only on-disk
+	// changes.
+	// TODO(rstambler): Make sure the go.work/gopls.mod files are always known
+	// to the view.
+	for _, src := range []workspaceSource{goWorkWorkspace, goplsModWorkspace} {
+		if c.URI == uriForSource(v.rootURI, src) {
+			return true
+		}
 	}
 	// If the file is not known to the view, and the change is only on-disk,
 	// we should not invalidate the snapshot. This is necessary because Emacs
@@ -775,7 +779,7 @@
 func findWorkspaceRoot(ctx context.Context, folder span.URI, fs source.FileSource, excludePath func(string) bool, experimental bool) (span.URI, error) {
 	patterns := []string{"go.mod"}
 	if experimental {
-		patterns = []string{"gopls.mod", "go.mod"}
+		patterns = []string{"go.work", "gopls.mod", "go.mod"}
 	}
 	for _, basename := range patterns {
 		dir, err := findRootPattern(ctx, folder, basename, fs)
diff --git a/internal/lsp/cache/workspace.go b/internal/lsp/cache/workspace.go
index d4b5303..4204bcc 100644
--- a/internal/lsp/cache/workspace.go
+++ b/internal/lsp/cache/workspace.go
@@ -6,6 +6,7 @@
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 	"sort"
@@ -15,6 +16,7 @@
 	"golang.org/x/mod/modfile"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/lsp/source"
+	workfile "golang.org/x/tools/internal/mod/modfile"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/xcontext"
 	errors "golang.org/x/xerrors"
@@ -25,6 +27,7 @@
 const (
 	legacyWorkspace = iota
 	goplsModWorkspace
+	goWorkWorkspace
 	fileSystemWorkspace
 )
 
@@ -34,6 +37,8 @@
 		return "legacy"
 	case goplsModWorkspace:
 		return "gopls.mod"
+	case goWorkWorkspace:
+		return "go.work"
 	case fileSystemWorkspace:
 		return "file system"
 	default:
@@ -86,24 +91,9 @@
 	// In experimental mode, the user may have a gopls.mod file that defines
 	// their workspace.
 	if experimental {
-		goplsModFH, err := fs.GetFile(ctx, goplsModURI(root))
-		if err != nil {
-			return nil, err
-		}
-		contents, err := goplsModFH.Read()
+		ws, err := parseExplicitWorkspaceFile(ctx, root, fs, excludePath)
 		if err == nil {
-			file, activeModFiles, err := parseGoplsMod(root, goplsModFH.URI(), contents)
-			if err != nil {
-				return nil, err
-			}
-			return &workspace{
-				root:           root,
-				excludePath:    excludePath,
-				activeModFiles: activeModFiles,
-				knownModFiles:  activeModFiles,
-				mod:            file,
-				moduleSource:   goplsModWorkspace,
-			}, nil
+			return ws, nil
 		}
 	}
 	// Otherwise, in all other modes, search for all of the go.mod files in the
@@ -145,6 +135,41 @@
 	}, nil
 }
 
+func parseExplicitWorkspaceFile(ctx context.Context, root span.URI, fs source.FileSource, excludePath func(string) bool) (*workspace, error) {
+	for _, src := range []workspaceSource{goWorkWorkspace, goplsModWorkspace} {
+		fh, err := fs.GetFile(ctx, uriForSource(root, src))
+		if err != nil {
+			return nil, err
+		}
+		contents, err := fh.Read()
+		if err != nil {
+			continue
+		}
+		var file *modfile.File
+		var activeModFiles map[span.URI]struct{}
+		switch src {
+		case goWorkWorkspace:
+			file, activeModFiles, err = parseGoWork(ctx, root, fh.URI(), contents, fs)
+		case goplsModWorkspace:
+			file, activeModFiles, err = parseGoplsMod(root, fh.URI(), contents)
+		}
+		if err != nil {
+			return nil, err
+		}
+		return &workspace{
+			root:           root,
+			excludePath:    excludePath,
+			activeModFiles: activeModFiles,
+			knownModFiles:  activeModFiles,
+			mod:            file,
+			moduleSource:   src,
+		}, nil
+	}
+	return nil, noHardcodedWorkspace
+}
+
+var noHardcodedWorkspace = errors.New("no hardcoded workspace")
+
 func (w *workspace) getKnownModFiles() map[span.URI]struct{} {
 	return w.knownModFiles
 }
@@ -246,7 +271,7 @@
 // Some workspace changes may affect workspace contents without requiring a
 // reload of metadata (for example, unsaved changes to a go.mod or go.sum
 // file).
-func (w *workspace) invalidate(ctx context.Context, changes map[span.URI]*fileChange) (_ *workspace, changed, reload bool) {
+func (w *workspace) invalidate(ctx context.Context, changes map[span.URI]*fileChange, fs source.FileSource) (_ *workspace, changed, reload bool) {
 	// Prevent races to w.modFile or w.wsDirs below, if wmhas not yet been built.
 	w.buildMu.Lock()
 	defer w.buildMu.Unlock()
@@ -269,42 +294,55 @@
 		result.activeModFiles[k] = v
 	}
 
-	// First handle changes to the gopls.mod file. This must be considered before
-	// any changes to go.mod or go.sum files, as the gopls.mod file determines
-	// which modules we care about. In legacy workspace mode we don't consider
-	// the gopls.mod file.
+	// First handle changes to the go.work or gopls.mod file. This must be
+	// considered before any changes to go.mod or go.sum files, as these files
+	// determine which modules we care about. In legacy workspace mode we don't
+	// consider the gopls.mod or go.work files.
 	if w.moduleSource != legacyWorkspace {
-		// If gopls.mod has changed we need to either re-read it if it exists or
-		// walk the filesystem if it has been deleted.
-		gmURI := goplsModURI(w.root)
-		// File opens/closes are just no-ops.
-		if change, ok := changes[gmURI]; ok && !change.isUnchanged {
+		// If go.work/gopls.mod has changed we need to either re-read it if it
+		// exists or walk the filesystem if it has been deleted.
+		// go.work should override the gopls.mod if both exist.
+		for _, src := range []workspaceSource{goplsModWorkspace, goWorkWorkspace} {
+			uri := uriForSource(w.root, src)
+			// File opens/closes are just no-ops.
+			change, ok := changes[uri]
+			if !ok || change.isUnchanged {
+				continue
+			}
 			if change.exists {
-				// Only invalidate if the gopls.mod actually parses.
-				// Otherwise, stick with the current gopls.mod.
-				parsedFile, parsedModules, err := parseGoplsMod(w.root, gmURI, change.content)
+				// Only invalidate if the file if it actually parses.
+				// Otherwise, stick with the current file.
+				var parsedFile *modfile.File
+				var parsedModules map[span.URI]struct{}
+				var err error
+				switch src {
+				case goWorkWorkspace:
+					parsedFile, parsedModules, err = parseGoWork(ctx, w.root, uri, change.content, fs)
+				case goplsModWorkspace:
+					parsedFile, parsedModules, err = parseGoplsMod(w.root, uri, change.content)
+				}
 				if err == nil {
 					changed = true
 					reload = change.fileHandle.Saved()
 					result.mod = parsedFile
-					result.moduleSource = goplsModWorkspace
+					result.moduleSource = src
 					result.knownModFiles = parsedModules
 					result.activeModFiles = make(map[span.URI]struct{})
 					for k, v := range parsedModules {
 						result.activeModFiles[k] = v
 					}
 				} else {
-					// An unparseable gopls.mod file should not invalidate the
-					// workspace: nothing good could come from changing the
-					// workspace in this case.
-					event.Error(ctx, "parsing gopls.mod", err)
+					// An unparseable file should not invalidate the workspace:
+					// nothing good could come from changing the workspace in
+					// this case.
+					event.Error(ctx, fmt.Sprintf("parsing %s", filepath.Base(uri.Filename())), err)
 				}
 			} else {
-				// gopls.mod is deleted. search for modules again.
+				// go.work/gopls.mod is deleted. search for modules again.
 				changed = true
 				reload = true
 				result.moduleSource = fileSystemWorkspace
-				// The parsed gopls.mod is no longer valid.
+				// The parsed file is no longer valid.
 				result.mod = nil
 				knownModFiles, err := findModules(w.root, w.excludePath, 0)
 				if err != nil {
@@ -325,7 +363,7 @@
 	// Next, handle go.mod changes that could affect our workspace. If we're
 	// reading our tracked modules from the gopls.mod, there's nothing to do
 	// here.
-	if result.moduleSource != goplsModWorkspace {
+	if result.moduleSource != goplsModWorkspace && result.moduleSource != goWorkWorkspace {
 		for uri, change := range changes {
 			if change.isUnchanged || !isGoMod(uri) || !source.InDir(result.root.Filename(), uri.Filename()) {
 				continue
@@ -370,8 +408,17 @@
 }
 
 // goplsModURI returns the URI for the gopls.mod file contained in root.
-func goplsModURI(root span.URI) span.URI {
-	return span.URIFromPath(filepath.Join(root.Filename(), "gopls.mod"))
+func uriForSource(root span.URI, src workspaceSource) span.URI {
+	var basename string
+	switch src {
+	case goplsModWorkspace:
+		basename = "gopls.mod"
+	case goWorkWorkspace:
+		basename = "go.work"
+	default:
+		return ""
+	}
+	return span.URIFromPath(filepath.Join(root.Filename(), basename))
 }
 
 // modURI returns the URI for the go.mod file contained in root.
@@ -429,6 +476,32 @@
 	return modules, nil
 }
 
+func parseGoWork(ctx context.Context, root, uri span.URI, contents []byte, fs source.FileSource) (*modfile.File, map[span.URI]struct{}, error) {
+	workFile, err := workfile.ParseWork(uri.Filename(), contents, nil)
+	if err != nil {
+		return nil, nil, errors.Errorf("parsing go.work: %w", err)
+	}
+	modFiles := make(map[span.URI]struct{})
+	for _, dir := range workFile.Directory {
+		// The resulting modfile must use absolute paths, so that it can be
+		// written to a temp directory.
+		dir.DiskPath = absolutePath(root, dir.DiskPath)
+		modURI := span.URIFromPath(filepath.Join(dir.DiskPath, "go.mod"))
+		modFiles[modURI] = struct{}{}
+	}
+	modFile, err := buildWorkspaceModFile(ctx, modFiles, fs)
+	if err != nil {
+		return nil, nil, err
+	}
+	if workFile.Go.Version != "" {
+		if err := modFile.AddGoStmt(workFile.Go.Version); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	return modFile, modFiles, nil
+}
+
 func parseGoplsMod(root, uri span.URI, contents []byte) (*modfile.File, map[span.URI]struct{}, error) {
 	modFile, err := modfile.Parse(uri.Filename(), contents, nil)
 	if err != nil {
@@ -439,19 +512,23 @@
 		if replace.New.Version != "" {
 			return nil, nil, errors.Errorf("gopls.mod: replaced module %q@%q must not have version", replace.New.Path, replace.New.Version)
 		}
-		dirFP := filepath.FromSlash(replace.New.Path)
-		if !filepath.IsAbs(dirFP) {
-			dirFP = filepath.Join(root.Filename(), dirFP)
-			// The resulting modfile must use absolute paths, so that it can be
-			// written to a temp directory.
-			replace.New.Path = dirFP
-		}
-		modURI := span.URIFromPath(filepath.Join(dirFP, "go.mod"))
+		// The resulting modfile must use absolute paths, so that it can be
+		// written to a temp directory.
+		replace.New.Path = absolutePath(root, replace.New.Path)
+		modURI := span.URIFromPath(filepath.Join(replace.New.Path, "go.mod"))
 		modFiles[modURI] = struct{}{}
 	}
 	return modFile, modFiles, nil
 }
 
+func absolutePath(root span.URI, path string) string {
+	dirFP := filepath.FromSlash(path)
+	if !filepath.IsAbs(dirFP) {
+		dirFP = filepath.Join(root.Filename(), dirFP)
+	}
+	return dirFP
+}
+
 // errExhausted is returned by findModules if the file scan limit is reached.
 var errExhausted = errors.New("exhausted")
 
diff --git a/internal/lsp/cache/workspace_test.go b/internal/lsp/cache/workspace_test.go
index 8524061..a03aedc 100644
--- a/internal/lsp/cache/workspace_test.go
+++ b/internal/lsp/cache/workspace_test.go
@@ -296,7 +296,7 @@
 						t.Fatal(err)
 					}
 				}
-				got, gotChanged, gotReload := w.invalidate(ctx, changes)
+				got, gotChanged, gotReload := w.invalidate(ctx, changes, fs)
 				if gotChanged != test.wantChanged {
 					t.Errorf("w.invalidate(): got changed %t, want %t", gotChanged, test.wantChanged)
 				}
diff --git a/internal/lsp/diagnostics.go b/internal/lsp/diagnostics.go
index 953c95c..07cc5ec 100644
--- a/internal/lsp/diagnostics.go
+++ b/internal/lsp/diagnostics.go
@@ -379,7 +379,7 @@
 		for _, d := range err.DiagList {
 			s.storeDiagnostics(snapshot, d.URI, modSource, []*source.Diagnostic{d})
 		}
-		errMsg = strings.Replace(err.MainError.Error(), "\n", " ", -1)
+		errMsg = strings.ReplaceAll(err.MainError.Error(), "\n", " ")
 	}
 
 	if s.criticalErrorStatus == nil {
diff --git a/internal/lsp/mod/diagnostics.go b/internal/lsp/mod/diagnostics.go
index 6495aeb..625bc63 100644
--- a/internal/lsp/mod/diagnostics.go
+++ b/internal/lsp/mod/diagnostics.go
@@ -88,7 +88,7 @@
 	// Packages in the workspace can contribute diagnostics to go.mod files.
 	wspkgs, err := snapshot.WorkspacePackages(ctx)
 	if err != nil && !source.IsNonFatalGoModError(err) {
-		event.Error(ctx, "diagnosing go.mod", err)
+		event.Error(ctx, fmt.Sprintf("workspace packages: diagnosing %s", pm.URI), err)
 	}
 	if err == nil {
 		for _, pkg := range wspkgs {
@@ -102,7 +102,7 @@
 
 	tidied, err := snapshot.ModTidy(ctx, pm)
 	if err != nil && !source.IsNonFatalGoModError(err) {
-		event.Error(ctx, "diagnosing go.mod", err)
+		event.Error(ctx, fmt.Sprintf("tidy: diagnosing %s", pm.URI), err)
 	}
 	if err == nil {
 		for _, d := range tidied.Diagnostics {