cmd/go: add workspace pruning mode

This change corrects a bug in the handling of module loading of
workspaces. Namely, there is an assumption by the module pruning code
that if a root module is selected then the packages of that module can
be resolved without loading the whole module graph. This is not true
in workspace mode because two workspace modules can require different
versions of a dependency. Worse, one workspace module can directly
require a depencency that is transitively required by another
workspace module, changing the version of that module loaded in the
fully expanded graph.

To correct this, a new 'workspace' pruning mode is added where the
roots are the workspace modules themselves, satisfying the assumption
made by the module pruning logic.

The rest of this change accounts for the new pruning mode where it's
used and correctly sets the requirements in this pruning mode.

Change-Id: I5d4d9877e492e196681f6ee9f8f18a08b4e95c61
Reviewed-on: https://go-review.googlesource.com/c/go/+/357169
Trust: Michael Matloob <matloob@golang.org>
Run-TryBot: Michael Matloob <matloob@golang.org>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>

src/cmd/go/internal/modload/buildlist.go

@@ -38,11 +38,17 @@ type Requirements struct {
 	// If pruned, the graph includes only the root modules, the explicit
 	// requirements of those root modules, and the transitive requirements of only
 	// the root modules that do not support pruning.
+	//
+	// If workspace, the graph includes only the workspace modules, the explicit
+	// requirements of the workspace modules, and the transitive requirements of
+	// the workspace modules that do not support pruning.
 	pruning modPruning
 
-	// rootModules is the set of module versions explicitly required by the main
-	// modules, sorted and capped to length. It may contain duplicates, and may
-	// contain multiple versions for a given module path.
+	// rootModules is the set of root modules of the graph, sorted and capped to
+	// length. It may contain duplicates, and may  contain multiple versions for a
+	// given module path. The root modules of the groph are the set of main
+	// modules in workspace mode, and the main module's direct requirements
+	// outside workspace mode.
 	rootModules    []module.Version
 	maxRootVersion map[string]string
 
@@ -99,6 +105,19 @@ var requirements *Requirements
 // If vendoring is in effect, the caller must invoke initVendor on the returned
 // *Requirements before any other method.
 func newRequirements(pruning modPruning, rootModules []module.Version, direct map[string]bool) *Requirements {
+	if pruning == workspace {
+		return &Requirements{
+			pruning:        pruning,
+			rootModules:    capVersionSlice(rootModules),
+			maxRootVersion: nil,
+			direct:         direct,
+		}
+	}
+
+	if inWorkspaceMode() && pruning != workspace {
+		panic("in workspace mode, but pruning is not workspace in newRequirements")
+	}
+
 	for i, m := range rootModules {
 		if m.Version == "" && MainModules.Contains(m.Path) {
 			panic(fmt.Sprintf("newRequirements called with untrimmed build list: rootModules[%v] is a main module", i))
@@ -291,13 +310,11 @@ func readModGraph(ctx context.Context, pruning modPruning, roots []module.Versio
 			g: mvs.NewGraph(cmpVersion, MainModules.Versions()),
 		}
 	)
-	for _, m := range MainModules.Versions() {
-		// Require all roots from all main modules.
-		_ = TODOWorkspaces("This flattens a level of the module graph, adding the dependencies " +
-			"of all main modules to a single requirements struct, and losing the information of which " +
-			"main module required which requirement. Rework the requirements struct and change this" +
-			"to reflect the structure of the main modules.")
-		mg.g.Require(m, roots)
+	if pruning != workspace {
+		if inWorkspaceMode() {
+			panic("pruning is not workspace in workspace mode")
+		}
+		mg.g.Require(MainModules.mustGetSingleMainModule(), roots)
 	}
 
 	var (
@@ -352,9 +369,13 @@ func readModGraph(ctx context.Context, pruning modPruning, roots []module.Versio
 			// are sufficient to build the packages it contains. We must load its full
 			// transitive dependency graph to be sure that we see all relevant
 			// dependencies.
-			if pruning == unpruned || summary.pruning == unpruned {
+			if pruning != pruned || summary.pruning == unpruned {
+				nextPruning := summary.pruning
+				if pruning == unpruned {
+					nextPruning = unpruned
+				}
 				for _, r := range summary.require {
-					enqueue(r, unpruned)
+					enqueue(r, nextPruning)
 				}
 			}
 		})
@@ -424,14 +445,17 @@ func (mg *ModuleGraph) findError() error {
 }
 
 func (mg *ModuleGraph) allRootsSelected() bool {
-	for _, mm := range MainModules.Versions() {
-		roots, _ := mg.g.RequiredBy(mm)
+	var roots []module.Version
+	if inWorkspaceMode() {
+		roots = MainModules.Versions()
+	} else {
+		roots, _ = mg.g.RequiredBy(MainModules.mustGetSingleMainModule())
+	}
 	for _, m := range roots {
 		if mg.Selected(m.Path) != m.Version {
 			return false
 		}
 	}
-	}
 	return true
 }
 
@@ -576,10 +600,29 @@ func tidyRoots(ctx context.Context, rs *Requirements, pkgs []*loadPkg) (*Require
 }
 
 func updateRoots(ctx context.Context, direct map[string]bool, rs *Requirements, pkgs []*loadPkg, add []module.Version, rootsImported bool) (*Requirements, error) {
-	if rs.pruning == unpruned {
+	switch rs.pruning {
+	case unpruned:
 		return updateUnprunedRoots(ctx, direct, rs, add)
-	}
+	case pruned:
 		return updatePrunedRoots(ctx, direct, rs, pkgs, add, rootsImported)
+	case workspace:
+		return updateWorkspaceRoots(ctx, rs, add)
+	default:
+		panic(fmt.Sprintf("unsupported pruning mode: %v", rs.pruning))
+	}
+}
+
+func updateWorkspaceRoots(ctx context.Context, rs *Requirements, add []module.Version) (*Requirements, error) {
+	if len(add) != 0 {
+		// add should be empty in workspace mode because a non-empty add slice means
+		// that there are missing roots in the current pruning mode or that the
+		// pruning mode is being changed. But the pruning mode should always be
+		// 'workspace' in workspace mode and the set of roots in workspace mode is
+		// always complete because it's the set of workspace modules, which can't
+		// be edited by loading.
+		panic("add is not empty")
+	}
+	return rs, nil
 }
 
 // tidyPrunedRoots returns a minimal set of root requirements that maintains the
@@ -1156,7 +1199,6 @@ func updateUnprunedRoots(ctx context.Context, direct map[string]bool, rs *Requir
 		}
 	}
 
-	// TODO(matloob): Make roots into a map.
 	var roots []module.Version
 	for _, mainModule := range MainModules.Versions() {
 		min, err := mvs.Req(mainModule, rootPaths, &mvsReqs{roots: keep})
@@ -1182,6 +1224,8 @@ func updateUnprunedRoots(ctx context.Context, direct map[string]bool, rs *Requir
 func convertPruning(ctx context.Context, rs *Requirements, pruning modPruning) (*Requirements, error) {
 	if rs.pruning == pruning {
 		return rs, nil
+	} else if rs.pruning == workspace || pruning == workspace {
+		panic("attempthing to convert to/from workspace pruning and another pruning type")
 	}
 
 	if pruning == unpruned {

src/cmd/go/internal/modload/init.go

@@ -704,7 +704,7 @@ func LoadModFile(ctx context.Context) *Requirements {
 		}
 	}
 
-	if MainModules.Index(mainModule).goVersionV == "" {
+	if MainModules.Index(mainModule).goVersionV == "" && rs.pruning != workspace {
 		// TODO(#45551): Do something more principled instead of checking
 		// cfg.CmdName directly here.
 		if cfg.BuildMod == "mod" && cfg.CmdName != "mod graph" && cfg.CmdName != "mod why" {
@@ -987,29 +987,29 @@ func makeMainModules(ms []module.Version, rootDirs []string, modFiles []*modfile
 // requirementsFromModFiles returns the set of non-excluded requirements from
 // the global modFile.
 func requirementsFromModFiles(ctx context.Context, modFiles []*modfile.File) *Requirements {
-	rootCap := 0
-	for i := range modFiles {
-		rootCap += len(modFiles[i].Require)
-	}
-	roots := make([]module.Version, 0, rootCap)
-	mPathCount := make(map[string]int)
-	for _, m := range MainModules.Versions() {
-		mPathCount[m.Path] = 1
-	}
+	var roots []module.Version
 	direct := map[string]bool{}
-	for _, modFile := range modFiles {
-	requirement:
+	var pruning modPruning
+	if inWorkspaceMode() {
+		pruning = workspace
+		roots = make([]module.Version, len(MainModules.Versions()))
+		copy(roots, MainModules.Versions())
+	} else {
+		pruning = pruningForGoVersion(MainModules.GoVersion())
+		if len(modFiles) != 1 {
+			panic(fmt.Errorf("requirementsFromModFiles called with %v modfiles outside workspace mode", len(modFiles)))
+		}
+		modFile := modFiles[0]
+		roots = make([]module.Version, 0, len(modFile.Require))
+		mm := MainModules.mustGetSingleMainModule()
 		for _, r := range modFile.Require {
-			// TODO(#45713): Maybe join
-			for _, mainModule := range MainModules.Versions() {
-				if index := MainModules.Index(mainModule); index != nil && index.exclude[r.Mod] {
+			if index := MainModules.Index(mm); index != nil && index.exclude[r.Mod] {
 				if cfg.BuildMod == "mod" {
 					fmt.Fprintf(os.Stderr, "go: dropping requirement on excluded version %s %s\n", r.Mod.Path, r.Mod.Version)
 				} else {
 					fmt.Fprintf(os.Stderr, "go: ignoring requirement on excluded version %s %s\n", r.Mod.Path, r.Mod.Version)
 				}
-					continue requirement
-				}
+				continue
 			}
 
 			roots = append(roots, r.Mod)
@@ -1019,7 +1019,7 @@ func requirementsFromModFiles(ctx context.Context, modFiles []*modfile.File) *Re
 		}
 	}
 	module.Sort(roots)
-	rs := newRequirements(pruningForGoVersion(MainModules.GoVersion()), roots, direct)
+	rs := newRequirements(pruning, roots, direct)
 	return rs
 }
 

src/cmd/go/internal/modload/load.go

@@ -1004,7 +1004,11 @@ func loadFromRoots(ctx context.Context, params loaderParams) *loader {
 	}
 
 	var err error
-	ld.requirements, err = convertPruning(ctx, ld.requirements, pruningForGoVersion(ld.GoVersion))
+	desiredPruning := pruningForGoVersion(ld.GoVersion)
+	if ld.requirements.pruning == workspace {
+		desiredPruning = workspace
+	}
+	ld.requirements, err = convertPruning(ctx, ld.requirements, desiredPruning)
 	if err != nil {
 		ld.errorf("go: %v\n", err)
 	}
@@ -1246,6 +1250,24 @@ func (ld *loader) updateRequirements(ctx context.Context) (changed bool, err err
 				continue
 			}
 
+			if inWorkspaceMode() {
+				// In workspace mode / workspace pruning mode, the roots are the main modules
+				// rather than the main module's direct dependencies. The check below on the selected
+				// roots does not apply.
+				if mg, err := rs.Graph(ctx); err != nil {
+					return false, err
+				} else if _, ok := mg.RequiredBy(dep.mod); !ok {
+					// dep.mod is not an explicit dependency, but needs to be.
+					// See comment on error returned below.
+					pkg.err = &DirectImportFromImplicitDependencyError{
+						ImporterPath: pkg.path,
+						ImportedPath: dep.path,
+						Module:       dep.mod,
+					}
+				}
+				continue
+			}
+
 			if pkg.err == nil && cfg.BuildMod != "mod" {
 				if v, ok := rs.rootSelected(dep.mod.Path); !ok || v != dep.mod.Version {
 					// dep.mod is not an explicit dependency, but needs to be.

src/cmd/go/internal/modload/modfile.go

@@ -120,6 +120,7 @@ type modPruning uint8
 const (
 	pruned    modPruning = iota // transitive dependencies of modules at go 1.17 and higher are pruned out
 	unpruned                    // no transitive dependencies are pruned out
+	workspace                   // pruned to the union of modules in the workspace
 )
 
 func pruningForGoVersion(goVersion string) modPruning {
@@ -554,7 +555,7 @@ type retraction struct {
 //
 // The caller must not modify the returned summary.
 func goModSummary(m module.Version) (*modFileSummary, error) {
-	if m.Version == "" && MainModules.Contains(m.Path) {
+	if m.Version == "" && !inWorkspaceMode() && MainModules.Contains(m.Path) {
 		panic("internal error: goModSummary called on a main module")
 	}
 
@@ -718,10 +719,15 @@ var rawGoModSummaryCache par.Cache // module.Version → rawGoModSummary result
 func rawGoModData(m module.Version) (name string, data []byte, err error) {
 	if m.Version == "" {
 		// m is a replacement module with only a file path.
+
 		dir := m.Path
 		if !filepath.IsAbs(dir) {
+			if inWorkspaceMode() && MainModules.Contains(m.Path) {
+				dir = MainModules.ModRoot(m)
+			} else {
 				dir = filepath.Join(replaceRelativeTo(), dir)
 			}
+		}
 		name = filepath.Join(dir, "go.mod")
 		if gomodActual, ok := fsys.OverlayPath(name); ok {
 			// Don't lock go.mod if it's part of the overlay.

src/cmd/go/testdata/script/work_prune.txt

@@ -14,7 +14,7 @@
 # TODO(#48331): We currently load the wrong version of q. Fix this.
 
 go list -m -f '{{.Version}}' example.com/q
-stdout '^v1.0.0$' # TODO(#48331): This should be 1.1.0. Fix this.
+stdout '^v1.1.0$'
 
 -- go.work --
 go 1.18