cmd/go/internal/modload: avoid loading the full module graph for imports satisfied by lazy roots

For #36460

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

5 files changed, 136 insertions, 87 deletions

diff --git a/src/cmd/go/internal/modload/import.go b/src/cmd/go/internal/modload/import.go
index 4e62e61bb0..6c863351ff 100644
--- a/src/cmd/go/internal/modload/import.go
+++ b/src/cmd/go/internal/modload/import.go
@@ -220,10 +220,13 @@ func (e *invalidImportError) Unwrap() error {
 	return e.err
 }
 
-// importFromModules finds the module and directory in the build list
-// containing the package with the given import path. The answer must be unique:
-// importFromModules returns an error if multiple modules attempt to provide
-// the same package.
+// importFromModules finds the module and directory in the dependency graph of
+// rs containing the package with the given import path. If mg is nil,
+// importFromModules attempts to locate the module using only the main module
+// and the roots of rs before it loads the full graph.
+//
+// The answer must be unique: importFromModules returns an error if multiple
+// modules are observed to provide the same package.
 //
 // importFromModules can return a module with an empty m.Path, for packages in
 // the standard library.
@@ -233,7 +236,7 @@ func (e *invalidImportError) Unwrap() error {
 //
 // If the package is not present in any module selected from the requirement
 // graph, importFromModules returns an *ImportMissingError.
-func importFromModules(ctx context.Context, path string, rs *Requirements) (m module.Version, dir string, err error) {
+func importFromModules(ctx context.Context, path string, rs *Requirements, mg *ModuleGraph) (m module.Version, dir string, err error) {
 	if strings.Contains(path, "@") {
 		return module.Version{}, "", fmt.Errorf("import path should not have @version")
 	}
@@ -295,93 +298,106 @@ func importFromModules(ctx context.Context, path string, rs *Requirements) (m mo
 	// large projects both M and P may be very large (note that M ≤ P), but k
 	// will tend to remain smallish (if for no other reason than filesystem
 	// path limitations).
-	var mg *ModuleGraph
-	if go117LazyTODO {
-		// Pull the prefix-matching loop below into another (new) loop.
-		// If the main module is lazy, try it once with mg == nil, and then load mg
-		// and try again.
-	} else {
-		mg, err = rs.Graph(ctx)
-		if err != nil {
-			// We might be missing one or more transitive (implicit) dependencies from
-			// the module graph, so we can't return an ImportMissingError here — one
-			// of the missing modules might actually contain the package in question,
-			// in which case we shouldn't go looking for it in some new dependency.
-			return module.Version{}, "", err
-		}
-	}
+	//
+	// We perform this iteration either one or two times. If mg is initially nil,
+	// then we first attempt to load the package using only the main module and
+	// its root requirements. If that does not identify the package, or if mg is
+	// already non-nil, then we attempt to load the package using the full
+	// requirements in mg.
+	for {
+		var sumErrMods []module.Version
+		for prefix := path; prefix != "."; prefix = pathpkg.Dir(prefix) {
+			var (
+				v  string
+				ok bool
+			)
+			if mg == nil {
+				v, ok = rs.rootSelected(prefix)
+			} else {
+				v, ok = mg.Selected(prefix), true
+			}
+			if !ok || v == "none" {
+				continue
+			}
+			m := module.Version{Path: prefix, Version: v}
 
-	var sumErrMods []module.Version
-	for prefix := path; prefix != "."; prefix = pathpkg.Dir(prefix) {
-		v := mg.Selected(prefix)
-		if v == "none" {
-			continue
+			needSum := true
+			root, isLocal, err := fetch(ctx, m, needSum)
+			if err != nil {
+				if sumErr := (*sumMissingError)(nil); errors.As(err, &sumErr) {
+					// We are missing a sum needed to fetch a module in the build list.
+					// We can't verify that the package is unique, and we may not find
+					// the package at all. Keep checking other modules to decide which
+					// error to report. Multiple sums may be missing if we need to look in
+					// multiple nested modules to resolve the import; we'll report them all.
+					sumErrMods = append(sumErrMods, m)
+					continue
+				}
+				// Report fetch error.
+				// Note that we don't know for sure this module is necessary,
+				// but it certainly _could_ provide the package, and even if we
+				// continue the loop and find the package in some other module,
+				// we need to look at this module to make sure the import is
+				// not ambiguous.
+				return module.Version{}, "", err
+			}
+			if dir, ok, err := dirInModule(path, m.Path, root, isLocal); err != nil {
+				return module.Version{}, "", err
+			} else if ok {
+				mods = append(mods, m)
+				dirs = append(dirs, dir)
+			}
 		}
-		m := module.Version{Path: prefix, Version: v}
 
-		needSum := true
-		root, isLocal, err := fetch(ctx, m, needSum)
-		if err != nil {
-			if sumErr := (*sumMissingError)(nil); errors.As(err, &sumErr) {
-				// We are missing a sum needed to fetch a module in the build list.
-				// We can't verify that the package is unique, and we may not find
-				// the package at all. Keep checking other modules to decide which
-				// error to report. Multiple sums may be missing if we need to look in
-				// multiple nested modules to resolve the import; we'll report them all.
-				sumErrMods = append(sumErrMods, m)
-				continue
+		if len(mods) > 1 {
+			// We produce the list of directories from longest to shortest candidate
+			// module path, but the AmbiguousImportError should report them from
+			// shortest to longest. Reverse them now.
+			for i := 0; i < len(mods)/2; i++ {
+				j := len(mods) - 1 - i
+				mods[i], mods[j] = mods[j], mods[i]
+				dirs[i], dirs[j] = dirs[j], dirs[i]
 			}
-			// Report fetch error.
-			// Note that we don't know for sure this module is necessary,
-			// but it certainly _could_ provide the package, and even if we
-			// continue the loop and find the package in some other module,
-			// we need to look at this module to make sure the import is
-			// not ambiguous.
-			return module.Version{}, "", err
+			return module.Version{}, "", &AmbiguousImportError{importPath: path, Dirs: dirs, Modules: mods}
 		}
-		if dir, ok, err := dirInModule(path, m.Path, root, isLocal); err != nil {
-			return module.Version{}, "", err
-		} else if ok {
-			mods = append(mods, m)
-			dirs = append(dirs, dir)
-		}
-	}
 
-	if len(mods) > 1 {
-		// We produce the list of directories from longest to shortest candidate
-		// module path, but the AmbiguousImportError should report them from
-		// shortest to longest. Reverse them now.
-		for i := 0; i < len(mods)/2; i++ {
-			j := len(mods) - 1 - i
-			mods[i], mods[j] = mods[j], mods[i]
-			dirs[i], dirs[j] = dirs[j], dirs[i]
+		if len(sumErrMods) > 0 {
+			for i := 0; i < len(sumErrMods)/2; i++ {
+				j := len(sumErrMods) - 1 - i
+				sumErrMods[i], sumErrMods[j] = sumErrMods[j], sumErrMods[i]
+			}
+			return module.Version{}, "", &ImportMissingSumError{
+				importPath: path,
+				mods:       sumErrMods,
+				found:      len(mods) > 0,
+			}
 		}
-		return module.Version{}, "", &AmbiguousImportError{importPath: path, Dirs: dirs, Modules: mods}
-	}
 
-	if len(sumErrMods) > 0 {
-		for i := 0; i < len(sumErrMods)/2; i++ {
-			j := len(sumErrMods) - 1 - i
-			sumErrMods[i], sumErrMods[j] = sumErrMods[j], sumErrMods[i]
+		if len(mods) == 1 {
+			return mods[0], dirs[0], nil
 		}
-		return module.Version{}, "", &ImportMissingSumError{
-			importPath: path,
-			mods:       sumErrMods,
-			found:      len(mods) > 0,
-		}
-	}
 
-	if len(mods) == 1 {
-		return mods[0], dirs[0], nil
-	}
+		if mg != nil {
+			// We checked the full module graph and still didn't find the
+			// requested package.
+			var queryErr error
+			if !HasModRoot() {
+				queryErr = ErrNoModRoot
+			}
+			return module.Version{}, "", &ImportMissingError{Path: path, QueryErr: queryErr, isStd: pathIsStd}
+		}
 
-	// We checked the full module graph and still didn't find the
-	// requested package.
-	var queryErr error
-	if !HasModRoot() {
-		queryErr = ErrNoModRoot
+		// So far we've checked the root dependencies.
+		// Load the full module graph and try again.
+		mg, err = rs.Graph(ctx)
+		if err != nil {
+			// We might be missing one or more transitive (implicit) dependencies from
+			// the module graph, so we can't return an ImportMissingError here — one
+			// of the missing modules might actually contain the package in question,
+			// in which case we shouldn't go looking for it in some new dependency.
+			return module.Version{}, "", err
+		}
 	}
-	return module.Version{}, "", &ImportMissingError{Path: path, QueryErr: queryErr, isStd: pathIsStd}
 }
 
 // queryImport attempts to locate a module that can be added to the current
diff --git a/src/cmd/go/internal/modload/init.go b/src/cmd/go/internal/modload/init.go
index cb206a3dea..f46c58f474 100644
--- a/src/cmd/go/internal/modload/init.go
+++ b/src/cmd/go/internal/modload/init.go
@@ -1125,6 +1125,22 @@ func keepSums(ctx context.Context, ld *loader, rs *Requirements, which whichSums
 				continue
 			}
 
+			if rs.depth == lazy && pkg.mod.Path != "" {
+				if v, ok := rs.rootSelected(pkg.mod.Path); ok && v == pkg.mod.Version {
+					// pkg was loaded from a root module, and because the main module is
+					// lazy we do not check non-root modules for conflicts for packages
+					// that can be found in roots. So we only need the checksums for the
+					// root modules that may contain pkg, not all possible modules.
+					for prefix := pkg.path; prefix != "."; prefix = path.Dir(prefix) {
+						if v, ok := rs.rootSelected(prefix); ok && v != "none" {
+							m := module.Version{Path: prefix, Version: v}
+							keep[resolveReplacement(m)] = true
+						}
+					}
+					continue
+				}
+			}
+
 			for prefix := pkg.path; prefix != "."; prefix = path.Dir(prefix) {
 				if v := mg.Selected(prefix); v != "none" {
 					m := module.Version{Path: prefix, Version: v}
diff --git a/src/cmd/go/internal/modload/load.go b/src/cmd/go/internal/modload/load.go
index b822e74eb5..ddacf49ead 100644
--- a/src/cmd/go/internal/modload/load.go
+++ b/src/cmd/go/internal/modload/load.go
@@ -1236,7 +1236,7 @@ func (ld *loader) updateRequirements(ctx context.Context) (changed bool, err err
 				//
 				// In some sense, we can think of this as ‘upgraded the module providing
 				// pkg.path from "none" to a version higher than "none"’.
-				if _, _, err = importFromModules(ctx, pkg.path, rs); err == nil {
+				if _, _, err = importFromModules(ctx, pkg.path, rs, nil); err == nil {
 					changed = true
 					break
 				}
@@ -1429,7 +1429,7 @@ func (ld *loader) preloadRootModules(ctx context.Context, rootPkgs []string) (ch
 			// If the main module is tidy and the package is in "all" — or if we're
 			// lucky — we can identify all of its imports without actually loading the
 			// full module graph.
-			m, _, err := importFromModules(ctx, path, ld.requirements)
+			m, _, err := importFromModules(ctx, path, ld.requirements, nil)
 			if err != nil {
 				var missing *ImportMissingError
 				if errors.As(err, &missing) && ld.ResolveMissingImports {
@@ -1516,7 +1516,24 @@ func (ld *loader) load(ctx context.Context, pkg *loadPkg) {
 		return
 	}
 
-	pkg.mod, pkg.dir, pkg.err = importFromModules(ctx, pkg.path, ld.requirements)
+	var mg *ModuleGraph
+	if ld.requirements.depth == eager {
+		var err error
+		mg, err = ld.requirements.Graph(ctx)
+		if err != nil {
+			// We already checked the error from Graph in loadFromRoots and/or
+			// updateRequirements, so we ignored the error on purpose and we should
+			// keep trying to push past it.
+			//
+			// However, because mg may be incomplete (and thus may select inaccurate
+			// versions), we shouldn't use it to load packages. Instead, we pass a nil
+			// *ModuleGraph, which will cause mg to first try loading from only the
+			// main module and root dependencies.
+			mg = nil
+		}
+	}
+
+	pkg.mod, pkg.dir, pkg.err = importFromModules(ctx, pkg.path, ld.requirements, mg)
 	if pkg.dir == "" {
 		return
 	}
diff --git a/src/cmd/go/testdata/script/mod_install_pkg_version.txt b/src/cmd/go/testdata/script/mod_install_pkg_version.txt
index b024fce174..2c14ef737b 100644
--- a/src/cmd/go/testdata/script/mod_install_pkg_version.txt
+++ b/src/cmd/go/testdata/script/mod_install_pkg_version.txt
@@ -67,9 +67,9 @@ cd tmp
 go mod init tmp
 go mod edit -require=rsc.io/fortune@v1.0.0
 ! go install -mod=readonly $GOPATH/pkg/mod/rsc.io/fortune@v1.0.0
-stderr '^rsc.io/fortune@v1.0.0: missing go.sum entry; to add it:\n\tgo mod download rsc.io/fortune$'
+stderr '^missing go\.sum entry for module providing package rsc\.io/fortune; to add:\n\tgo mod download rsc\.io/fortune$'
 ! go install -mod=readonly ../../pkg/mod/rsc.io/fortune@v1.0.0
-stderr '^rsc.io/fortune@v1.0.0: missing go.sum entry; to add it:\n\tgo mod download rsc.io/fortune$'
+stderr '^missing go\.sum entry for module providing package rsc\.io/fortune; to add:\n\tgo mod download rsc\.io/fortune$'
 go get -d rsc.io/fortune@v1.0.0
 go install -mod=readonly $GOPATH/pkg/mod/rsc.io/fortune@v1.0.0
 exists $GOPATH/bin/fortune$GOEXE
diff --git a/src/cmd/go/testdata/script/mod_run_pkg_version.txt b/src/cmd/go/testdata/script/mod_run_pkg_version.txt
index 3c3ed27e91..e921fab508 100644
--- a/src/cmd/go/testdata/script/mod_run_pkg_version.txt
+++ b/src/cmd/go/testdata/script/mod_run_pkg_version.txt
@@ -64,9 +64,9 @@ cd tmp
 go mod init tmp
 go mod edit -require=rsc.io/fortune@v1.0.0
 ! go run -mod=readonly $GOPATH/pkg/mod/rsc.io/fortune@v1.0.0
-stderr '^rsc.io/fortune@v1.0.0: missing go.sum entry; to add it:\n\tgo mod download rsc.io/fortune$'
+stderr '^missing go\.sum entry for module providing package rsc\.io/fortune; to add:\n\tgo mod download rsc\.io/fortune$'
 ! go run -mod=readonly ../../pkg/mod/rsc.io/fortune@v1.0.0
-stderr '^rsc.io/fortune@v1.0.0: missing go.sum entry; to add it:\n\tgo mod download rsc.io/fortune$'
+stderr '^missing go\.sum entry for module providing package rsc\.io/fortune; to add:\n\tgo mod download rsc\.io/fortune$'
 cd ..
 rm tmp