8 files changed, 363 insertions, 7 deletions

diff --git a/src/cmd/go/internal/modcmd/download.go b/src/cmd/go/internal/modcmd/download.go
index 0e5af85237..ff56d05116 100644
--- a/src/cmd/go/internal/modcmd/download.go
+++ b/src/cmd/go/internal/modcmd/download.go
@@ -66,6 +66,7 @@ func init() {
 	// TODO(jayconrod): https://golang.org/issue/35849 Apply -x to other 'go mod' commands.
 	cmdDownload.Flag.BoolVar(&cfg.BuildX, "x", false, "")
 	base.AddModCommonFlags(&cmdDownload.Flag)
+	base.AddWorkfileFlag(&cmdDownload.Flag)
 }
 
 type moduleJSON struct {
@@ -81,6 +82,8 @@ type moduleJSON struct {
 }
 
 func runDownload(ctx context.Context, cmd *base.Command, args []string) {
+	modload.InitWorkfile()
+
 	// Check whether modules are enabled and whether we're in a module.
 	modload.ForceUseModules = true
 	if !modload.HasModRoot() && len(args) == 0 {
@@ -91,12 +94,18 @@ func runDownload(ctx context.Context, cmd *base.Command, args []string) {
 		args = []string{"all"}
 	}
 	if modload.HasModRoot() {
-		modload.LoadModFile(ctx) // to fill Target
-		targetAtUpgrade := modload.Target.Path + "@upgrade"
-		targetAtPatch := modload.Target.Path + "@patch"
+		modload.LoadModFile(ctx) // to fill MainModules
+
+		if len(modload.MainModules.Versions()) != 1 {
+			panic(modload.TODOWorkspaces("Support workspace mode in go mod download"))
+		}
+		mainModule := modload.MainModules.Versions()[0]
+
+		targetAtUpgrade := mainModule.Path + "@upgrade"
+		targetAtPatch := mainModule.Path + "@patch"
 		for _, arg := range args {
 			switch arg {
-			case modload.Target.Path, targetAtUpgrade, targetAtPatch:
+			case mainModule.Path, targetAtUpgrade, targetAtPatch:
 				os.Stderr.WriteString("go mod download: skipping argument " + arg + " that resolves to the main module\n")
 			}
 		}
diff --git a/src/cmd/go/internal/modcmd/editwork.go b/src/cmd/go/internal/modcmd/editwork.go
new file mode 100644
index 0000000000..f05d9245e7
--- /dev/null
+++ b/src/cmd/go/internal/modcmd/editwork.go
@@ -0,0 +1,282 @@
+// Copyright 2021 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.
+
+// go mod editwork
+
+package modcmd
+
+import (
+	"bytes"
+	"cmd/go/internal/base"
+	"cmd/go/internal/lockedfile"
+	"cmd/go/internal/modload"
+	"context"
+	"encoding/json"
+	"errors"
+	"os"
+	"strings"
+
+	"golang.org/x/mod/modfile"
+)
+
+var cmdEditwork = &base.Command{
+	UsageLine: "go mod editwork [editing flags] [go.work]",
+	Short:     "edit go.work from tools or scripts",
+	Long: `Editwork provides a command-line interface for editing go.work,
+for use primarily by tools or scripts. It only reads go.work;
+it does not look up information about the modules involved.
+If no file is specified, editwork looks for a go.work file in the current
+directory and its parent directories
+
+The editing flags specify a sequence of editing operations.
+
+The -fmt flag reformats the go.work file without making other changes.
+This reformatting is also implied by any other modifications that use or
+rewrite the go.mod file. The only time this flag is needed is if no other
+flags are specified, as in 'go mod editwork -fmt'.
+
+The -directory=path and -dropdirectory=path flags
+add and drop a directory from the go.work files set of module directories.
+
+The -replace=old[@v]=new[@v] flag adds a replacement of the given
+module path and version pair. If the @v in old@v is omitted, a
+replacement without a version on the left side is added, which applies
+to all versions of the old module path. If the @v in new@v is omitted,
+the new path should be a local module root directory, not a module
+path. Note that -replace overrides any redundant replacements for old[@v],
+so omitting @v will drop existing replacements for specific versions.
+
+The -dropreplace=old[@v] flag drops a replacement of the given
+module path and version pair. If the @v is omitted, a replacement without
+a version on the left side is dropped.
+
+The -directory, -dropdirectory, -replace, and -dropreplace,
+editing flags may be repeated, and the changes are applied in the order given.
+
+The -go=version flag sets the expected Go language version.
+
+The -print flag prints the final go.work in its text format instead of
+writing it back to go.mod.
+
+The -json flag prints the final go.work file in JSON format instead of
+writing it back to go.mod. The JSON output corresponds to these Go types:
+
+	type Module struct {
+		Path    string
+		Version string
+	}
+
+	type GoWork struct {
+		Go        string
+		Directory []Directory
+		Replace   []Replace
+	}
+
+	type Directory struct {
+		Path       string
+		ModulePath string
+	}
+
+	type Replace struct {
+		Old Module
+		New Module
+	}
+
+See the workspaces design proposal at
+https://go.googlesource.com/proposal/+/master/design/45713-workspace.md for
+more information.
+`,
+}
+
+var (
+	editworkFmt   = cmdEditwork.Flag.Bool("fmt", false, "")
+	editworkGo    = cmdEditwork.Flag.String("go", "", "")
+	editworkJSON  = cmdEditwork.Flag.Bool("json", false, "")
+	editworkPrint = cmdEditwork.Flag.Bool("print", false, "")
+	workedits     []func(file *modfile.WorkFile) // edits specified in flags
+)
+
+func init() {
+	cmdEditwork.Run = runEditwork // break init cycle
+
+	cmdEditwork.Flag.Var(flagFunc(flagEditworkDirectory), "directory", "")
+	cmdEditwork.Flag.Var(flagFunc(flagEditworkDropDirectory), "dropdirectory", "")
+	cmdEditwork.Flag.Var(flagFunc(flagEditworkReplace), "replace", "")
+	cmdEditwork.Flag.Var(flagFunc(flagEditworkDropReplace), "dropreplace", "")
+
+	base.AddWorkfileFlag(&cmdEditwork.Flag)
+}
+
+func runEditwork(ctx context.Context, cmd *base.Command, args []string) {
+	anyFlags :=
+		*editworkGo != "" ||
+			*editworkJSON ||
+			*editworkPrint ||
+			*editworkFmt ||
+			len(workedits) > 0
+
+	if !anyFlags {
+		base.Fatalf("go mod edit: no flags specified (see 'go help mod edit').")
+	}
+
+	if *editworkJSON && *editworkPrint {
+		base.Fatalf("go mod edit: cannot use both -json and -print")
+	}
+
+	if len(args) > 1 {
+		base.Fatalf("go mod edit: too many arguments")
+	}
+	var gowork string
+	if len(args) == 1 {
+		gowork = args[0]
+	} else {
+		modload.InitWorkfile()
+		gowork = modload.WorkFilePath()
+	}
+
+	if *editworkGo != "" {
+		if !modfile.GoVersionRE.MatchString(*editworkGo) {
+			base.Fatalf(`go mod: invalid -go option; expecting something like "-go %s"`, modload.LatestGoVersion())
+		}
+	}
+
+	data, err := lockedfile.Read(gowork)
+	if err != nil {
+		base.Fatalf("go: %v", err)
+	}
+
+	workFile, err := modfile.ParseWork(gowork, data, nil)
+	if err != nil {
+		base.Fatalf("go: errors parsing %s:\n%s", base.ShortPath(gowork), err)
+	}
+
+	if *editworkGo != "" {
+		if err := workFile.AddGoStmt(*editworkGo); err != nil {
+			base.Fatalf("go: internal error: %v", err)
+		}
+	}
+
+	if len(workedits) > 0 {
+		for _, edit := range workedits {
+			edit(workFile)
+		}
+	}
+	workFile.SortBlocks()
+	workFile.Cleanup() // clean file after edits
+
+	if *editworkJSON {
+		editworkPrintJSON(workFile)
+		return
+	}
+
+	out := modfile.Format(workFile.Syntax)
+
+	if *editworkPrint {
+		os.Stdout.Write(out)
+		return
+	}
+
+	err = lockedfile.Transform(gowork, func(lockedData []byte) ([]byte, error) {
+		if !bytes.Equal(lockedData, data) {
+			return nil, errors.New("go.work changed during editing; not overwriting")
+		}
+		return out, nil
+	})
+	if err != nil {
+		base.Fatalf("go: %v", err)
+	}
+}
+
+// flagEditworkDirectory implements the -directory flag.
+func flagEditworkDirectory(arg string) {
+	workedits = append(workedits, func(f *modfile.WorkFile) {
+		if err := f.AddDirectory(arg, ""); err != nil {
+			base.Fatalf("go mod: -directory=%s: %v", arg, err)
+		}
+	})
+}
+
+// flagEditworkDropDirectory implements the -dropdirectory flag.
+func flagEditworkDropDirectory(arg string) {
+	workedits = append(workedits, func(f *modfile.WorkFile) {
+		if err := f.DropDirectory(arg); err != nil {
+			base.Fatalf("go mod: -dropdirectory=%s: %v", arg, err)
+		}
+	})
+}
+
+// flagReplace implements the -replace flag.
+func flagEditworkReplace(arg string) {
+	var i int
+	if i = strings.Index(arg, "="); i < 0 {
+		base.Fatalf("go mod: -replace=%s: need old[@v]=new[@w] (missing =)", arg)
+	}
+	old, new := strings.TrimSpace(arg[:i]), strings.TrimSpace(arg[i+1:])
+	if strings.HasPrefix(new, ">") {
+		base.Fatalf("go mod: -replace=%s: separator between old and new is =, not =>", arg)
+	}
+	oldPath, oldVersion, err := parsePathVersionOptional("old", old, false)
+	if err != nil {
+		base.Fatalf("go mod: -replace=%s: %v", arg, err)
+	}
+	newPath, newVersion, err := parsePathVersionOptional("new", new, true)
+	if err != nil {
+		base.Fatalf("go mod: -replace=%s: %v", arg, err)
+	}
+	if newPath == new && !modfile.IsDirectoryPath(new) {
+		base.Fatalf("go mod: -replace=%s: unversioned new path must be local directory", arg)
+	}
+
+	workedits = append(workedits, func(f *modfile.WorkFile) {
+		if err := f.AddReplace(oldPath, oldVersion, newPath, newVersion); err != nil {
+			base.Fatalf("go mod: -replace=%s: %v", arg, err)
+		}
+	})
+}
+
+// flagDropReplace implements the -dropreplace flag.
+func flagEditworkDropReplace(arg string) {
+	path, version, err := parsePathVersionOptional("old", arg, true)
+	if err != nil {
+		base.Fatalf("go mod: -dropreplace=%s: %v", arg, err)
+	}
+	workedits = append(workedits, func(f *modfile.WorkFile) {
+		if err := f.DropReplace(path, version); err != nil {
+			base.Fatalf("go mod: -dropreplace=%s: %v", arg, err)
+		}
+	})
+}
+
+// editPrintJSON prints the -json output.
+func editworkPrintJSON(workFile *modfile.WorkFile) {
+	var f workfileJSON
+	if workFile.Go != nil {
+		f.Go = workFile.Go.Version
+	}
+	for _, d := range workFile.Directory {
+		f.Directory = append(f.Directory, directoryJSON{DiskPath: d.Path, ModPath: d.ModulePath})
+	}
+
+	for _, r := range workFile.Replace {
+		f.Replace = append(f.Replace, replaceJSON{r.Old, r.New})
+	}
+	data, err := json.MarshalIndent(&f, "", "\t")
+	if err != nil {
+		base.Fatalf("go: internal error: %v", err)
+	}
+	data = append(data, '\n')
+	os.Stdout.Write(data)
+}
+
+// workfileJSON is the -json output data structure.
+type workfileJSON struct {
+	Go        string `json:",omitempty"`
+	Directory []directoryJSON
+	Replace   []replaceJSON
+}
+
+type directoryJSON struct {
+	DiskPath string
+	ModPath  string `json:",omitempty"`
+}
diff --git a/src/cmd/go/internal/modcmd/graph.go b/src/cmd/go/internal/modcmd/graph.go
index ac81f26dad..2cbabae044 100644
--- a/src/cmd/go/internal/modcmd/graph.go
+++ b/src/cmd/go/internal/modcmd/graph.go
@@ -42,9 +42,12 @@ var (
 func init() {
 	cmdGraph.Flag.Var(&graphGo, "go", "")
 	base.AddModCommonFlags(&cmdGraph.Flag)
+	base.AddWorkfileFlag(&cmdGraph.Flag)
 }
 
 func runGraph(ctx context.Context, cmd *base.Command, args []string) {
+	modload.InitWorkfile()
+
 	if len(args) > 0 {
 		base.Fatalf("go mod graph: graph takes no arguments")
 	}
diff --git a/src/cmd/go/internal/modcmd/initwork.go b/src/cmd/go/internal/modcmd/initwork.go
new file mode 100644
index 0000000000..4182aa071d
--- /dev/null
+++ b/src/cmd/go/internal/modcmd/initwork.go
@@ -0,0 +1,54 @@
+// Copyright 2021 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.
+
+// go mod initwork
+
+package modcmd
+
+import (
+	"cmd/go/internal/base"
+	"cmd/go/internal/modload"
+	"context"
+	"path/filepath"
+)
+
+var _ = modload.TODOWorkspaces("Add more documentation below. Though this is" +
+	"enough for those trying workspaces out, there should be more through" +
+	"documentation if the proposal is accepted and released.")
+
+var cmdInitwork = &base.Command{
+	UsageLine: "go mod initwork [moddirs]",
+	Short:     "initialize workspace file",
+	Long: `go mod initwork initializes and writes a new go.work file in the current
+directory, in effect creating a new workspace at the current directory.
+
+go mod initwork optionally accepts paths to the workspace modules as arguments.
+If the argument is omitted, an empty workspace with no modules will be created.
+
+See the workspaces design proposal at
+https://go.googlesource.com/proposal/+/master/design/45713-workspace.md for
+more information.
+`,
+	Run: runInitwork,
+}
+
+func init() {
+	base.AddModCommonFlags(&cmdInitwork.Flag)
+	base.AddWorkfileFlag(&cmdInitwork.Flag)
+}
+
+func runInitwork(ctx context.Context, cmd *base.Command, args []string) {
+	modload.InitWorkfile()
+
+	modload.ForceUseModules = true
+
+	// TODO(matloob): support using the -workfile path
+	// To do that properly, we'll have to make the module directories
+	// make dirs relative to workFile path before adding the paths to
+	// the directory entries
+
+	workFile := filepath.Join(base.Cwd(), "go.work")
+
+	modload.CreateWorkFile(ctx, workFile, args)
+}
diff --git a/src/cmd/go/internal/modcmd/mod.go b/src/cmd/go/internal/modcmd/mod.go
index d72d0cacd6..29aad58324 100644
--- a/src/cmd/go/internal/modcmd/mod.go
+++ b/src/cmd/go/internal/modcmd/mod.go
@@ -23,8 +23,10 @@ See 'go help modules' for an overview of module functionality.
 	Commands: []*base.Command{
 		cmdDownload,
 		cmdEdit,
+		cmdEditwork,
 		cmdGraph,
 		cmdInit,
+		cmdInitwork,
 		cmdTidy,
 		cmdVendor,
 		cmdVerify,
diff --git a/src/cmd/go/internal/modcmd/vendor.go b/src/cmd/go/internal/modcmd/vendor.go
index 4e01dfab1e..1effcea1a0 100644
--- a/src/cmd/go/internal/modcmd/vendor.go
+++ b/src/cmd/go/internal/modcmd/vendor.go
@@ -74,7 +74,7 @@ func runVendor(ctx context.Context, cmd *base.Command, args []string) {
 	}
 	_, pkgs := modload.LoadPackages(ctx, loadOpts, "all")
 
-	vdir := filepath.Join(modload.ModRoot(), "vendor")
+	vdir := filepath.Join(modload.VendorDir())
 	if err := os.RemoveAll(vdir); err != nil {
 		base.Fatalf("go mod vendor: %v", err)
 	}
@@ -82,7 +82,7 @@ func runVendor(ctx context.Context, cmd *base.Command, args []string) {
 	modpkgs := make(map[module.Version][]string)
 	for _, pkg := range pkgs {
 		m := modload.PackageModule(pkg)
-		if m.Path == "" || m == modload.Target {
+		if m.Path == "" || m.Version == "" && modload.MainModules.Contains(m.Path) {
 			continue
 		}
 		modpkgs[m] = append(modpkgs[m], pkg)
@@ -128,7 +128,8 @@ func runVendor(ctx context.Context, cmd *base.Command, args []string) {
 	}
 
 	for _, m := range vendorMods {
-		line := moduleLine(m, modload.Replacement(m))
+		replacement, _ := modload.Replacement(m)
+		line := moduleLine(m, replacement)
 		io.WriteString(w, line)
 
 		goVersion := ""
diff --git a/src/cmd/go/internal/modcmd/verify.go b/src/cmd/go/internal/modcmd/verify.go
index 5a6eca32cf..14c4d76bc3 100644
--- a/src/cmd/go/internal/modcmd/verify.go
+++ b/src/cmd/go/internal/modcmd/verify.go
@@ -39,9 +39,12 @@ See https://golang.org/ref/mod#go-mod-verify for more about 'go mod verify'.
 
 func init() {
 	base.AddModCommonFlags(&cmdVerify.Flag)
+	base.AddWorkfileFlag(&cmdVerify.Flag)
 }
 
 func runVerify(ctx context.Context, cmd *base.Command, args []string) {
+	modload.InitWorkfile()
+
 	if len(args) != 0 {
 		// NOTE(rsc): Could take a module pattern.
 		base.Fatalf("go mod verify: verify takes no arguments")
diff --git a/src/cmd/go/internal/modcmd/why.go b/src/cmd/go/internal/modcmd/why.go
index 3b14b27c8c..eef5fa5ae8 100644
--- a/src/cmd/go/internal/modcmd/why.go
+++ b/src/cmd/go/internal/modcmd/why.go
@@ -61,9 +61,11 @@ var (
 func init() {
 	cmdWhy.Run = runWhy // break init cycle
 	base.AddModCommonFlags(&cmdWhy.Flag)
+	base.AddWorkfileFlag(&cmdWhy.Flag)
 }
 
 func runWhy(ctx context.Context, cmd *base.Command, args []string) {
+	modload.InitWorkfile()
 	modload.ForceUseModules = true
 	modload.RootMode = modload.NeedRoot