go/packages: add Dir and ForTest fields to Package

Add a new Dir field, and export ForTest, per golang/go#38445. The
needInternalForTest mode bit is promoted to NeedForTest, but no mode bit
is added for Dir as it is logically related to NeedFiles.

A test is added for the new fields using a simpler txtar-based setup,
since I did not have time to page in the quite heavyweight packagestest
framework.

For golang/go#38445

Change-Id: I92026462f7ed7e237db1f4e50a3bbf2936802fbb
Reviewed-on: https://go-review.googlesource.com/c/tools/+/626495
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

go/packages/golist.go

@@ -505,13 +505,14 @@ func (state *golistState) createDriverResponse(words ...string) (*DriverResponse
 		pkg := &Package{
 			Name:            p.Name,
 			ID:              p.ImportPath,
+			Dir:             p.Dir,
 			GoFiles:         absJoin(p.Dir, p.GoFiles, p.CgoFiles),
 			CompiledGoFiles: absJoin(p.Dir, p.CompiledGoFiles),
 			OtherFiles:      absJoin(p.Dir, otherFiles(p)...),
 			EmbedFiles:      absJoin(p.Dir, p.EmbedFiles),
 			EmbedPatterns:   absJoin(p.Dir, p.EmbedPatterns),
 			IgnoredFiles:    absJoin(p.Dir, p.IgnoredGoFiles, p.IgnoredOtherFiles),
-			forTest:         p.ForTest,
+			ForTest:         p.ForTest,
 			depsErrors:      p.DepsErrors,
 			Module:          p.Module,
 		}
@@ -795,7 +796,7 @@ func jsonFlag(cfg *Config, goVersion int) string {
 		// Request Dir in the unlikely case Export is not absolute.
 		addFields("Dir", "Export")
 	}
-	if cfg.Mode&needInternalForTest != 0 {
+	if cfg.Mode&NeedForTest != 0 {
 		addFields("ForTest")
 	}
 	if cfg.Mode&needInternalDepsErrors != 0 {

go/packages/loadmode_string.go

@@ -23,6 +23,7 @@ var modes = [...]struct {
 	{NeedSyntax, "NeedSyntax"},
 	{NeedTypesInfo, "NeedTypesInfo"},
 	{NeedTypesSizes, "NeedTypesSizes"},
+	{NeedForTest, "NeedForTest"},
 	{NeedModule, "NeedModule"},
 	{NeedEmbedFiles, "NeedEmbedFiles"},
 	{NeedEmbedPatterns, "NeedEmbedPatterns"},

go/packages/packages.go

@@ -55,7 +55,7 @@ const (
 	// NeedName adds Name and PkgPath.
 	NeedName LoadMode = 1 << iota
 
-	// NeedFiles adds GoFiles, OtherFiles, and IgnoredFiles
+	// NeedFiles adds Dir, GoFiles, OtherFiles, and IgnoredFiles
 	NeedFiles
 
 	// NeedCompiledGoFiles adds CompiledGoFiles.
@@ -86,9 +86,10 @@ const (
 	// needInternalDepsErrors adds the internal deps errors field for use by gopls.
 	needInternalDepsErrors
 
-	// needInternalForTest adds the internal forTest field.
+	// NeedForTest adds ForTest.
+	//
 	// Tests must also be set on the context for this field to be populated.
-	needInternalForTest
+	NeedForTest
 
 	// typecheckCgo enables full support for type checking cgo. Requires Go 1.15+.
 	// Modifies CompiledGoFiles and Types, and has no effect on its own.
@@ -434,6 +435,12 @@ type Package struct {
 	// PkgPath is the package path as used by the go/types package.
 	PkgPath string
 
+	// Dir is the directory associated with the package, if it exists.
+	//
+	// For packages listed by the go command, this is the directory containing
+	// the package files.
+	Dir string
+
 	// Errors contains any errors encountered querying the metadata
 	// of the package, or while parsing or type-checking its files.
 	Errors []Error
@@ -521,8 +528,8 @@ type Package struct {
 
 	// -- internal --
 
-	// forTest is the package under test, if any.
-	forTest string
+	// ForTest is the package under test, if any.
+	ForTest string
 
 	// depsErrors is the DepsErrors field from the go list response, if any.
 	depsErrors []*packagesinternal.PackageError
@@ -551,9 +558,6 @@ type ModuleError struct {
 }
 
 func init() {
-	packagesinternal.GetForTest = func(p interface{}) string {
-		return p.(*Package).forTest
-	}
 	packagesinternal.GetDepsErrors = func(p interface{}) []*packagesinternal.PackageError {
 		return p.(*Package).depsErrors
 	}
@@ -565,7 +569,6 @@ func init() {
 	}
 	packagesinternal.TypecheckCgo = int(typecheckCgo)
 	packagesinternal.DepsErrors = int(needInternalDepsErrors)
-	packagesinternal.ForTest = int(needInternalForTest)
 }
 
 // An Error describes a problem with a package's metadata, syntax, or types.

go/packages/packages_test.go

@@ -26,11 +26,13 @@ import (
 	"testing/fstest"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/packagesinternal"
 	"golang.org/x/tools/internal/packagestest"
 	"golang.org/x/tools/internal/testenv"
 	"golang.org/x/tools/internal/testfiles"
+	"golang.org/x/tools/txtar"
 )
 
 // testCtx is canceled when the test binary is about to time out.
@@ -2324,6 +2326,10 @@ func TestLoadModeStrings(t *testing.T) {
 			packages.NeedName | packages.NeedExportFile,
 			"(NeedName|NeedExportFile)",
 		},
+		{
+			packages.NeedForTest | packages.NeedEmbedFiles | packages.NeedEmbedPatterns,
+			"(NeedForTest|NeedEmbedFiles|NeedEmbedPatterns)",
+		},
 		{
 			packages.NeedName | packages.NeedFiles | packages.NeedCompiledGoFiles | packages.NeedImports | packages.NeedDeps | packages.NeedExportFile | packages.NeedTypes | packages.NeedSyntax | packages.NeedTypesInfo | packages.NeedTypesSizes,
 			"(NeedName|NeedFiles|NeedCompiledGoFiles|NeedImports|NeedDeps|NeedExportFile|NeedTypes|NeedSyntax|NeedTypesInfo|NeedTypesSizes)",
@@ -2425,8 +2431,7 @@ func testForTestField(t *testing.T, exporter packagestest.Exporter) {
 		if !hasTestFile {
 			continue
 		}
-		got := packagesinternal.GetForTest(pkg)
-		if got != forTest {
+		if got := pkg.ForTest; got != forTest {
 			t.Errorf("expected %q, got %q", forTest, got)
 		}
 	}
@@ -3209,3 +3214,102 @@ func foo() int {
 		t.Errorf("expected types info to be present but got nil")
 	}
 }
+
+// TestDirAndForTest tests the new fields added as part of golang/go#38445.
+func TestDirAndForTest(t *testing.T) {
+	dir := writeTree(t, `
+-- go.mod --
+module example.com
+
+go 1.18
+
+-- a/a.go --
+package a
+
+func Foo() int { return 1 }
+
+-- a/a_test.go --
+package a
+
+func Bar() int { return 2 }
+
+-- a/a_x_test.go --
+package a_test
+
+import (
+	"testing"
+
+	"example.com/a"
+	"example.com/b"
+)
+
+func TestFooBar(t *testing.T) {
+	if got := a.Foo() + a.Bar() + b.Baz(); got != 6 {
+		t.Errorf("whoops!")
+	}
+}
+
+-- b/b.go --
+package b
+
+import "example.com/a"
+
+func Baz() int { return 3 }
+
+func Foo() int { return a.Foo() }
+`)
+
+	pkgs, err := packages.Load(&packages.Config{
+		Mode: packages.NeedName |
+			packages.NeedFiles |
+			packages.NeedForTest |
+			packages.NeedImports,
+		Dir:   dir,
+		Tests: true,
+	}, "./...")
+	if err != nil {
+		t.Fatal(err)
+	}
+	type result struct{ Dir, ForTest string }
+	got := make(map[string]result)
+	packages.Visit(pkgs, nil, func(pkg *packages.Package) {
+		if strings.Contains(pkg.PkgPath, ".") { // ignore std
+			rel, err := filepath.Rel(dir, pkg.Dir)
+			if err != nil {
+				t.Errorf("Rel(%q, %q) failed: %v", dir, pkg.Dir, err)
+				return
+			}
+			got[pkg.ID] = result{
+				Dir:     rel,
+				ForTest: pkg.ForTest,
+			}
+		}
+	})
+	want := map[string]result{
+		"example.com/a":                           {"a", ""},
+		"example.com/a.test":                      {"a", ""},
+		"example.com/a [example.com/a.test]":      {"a", "example.com/a"}, // test variant
+		"example.com/a_test [example.com/a.test]": {"a", "example.com/a"}, // x_test
+		"example.com/b [example.com/a.test]":      {"b", "example.com/a"}, // intermediate test variant
+		"example.com/b":                           {"b", ""},
+	}
+	if diff := cmp.Diff(want, got); diff != "" {
+		t.Errorf("Load returned mismatching ForTest fields (ID->result -want +got):\n%s", diff)
+	}
+	t.Logf("Packages: %+v", pkgs)
+}
+
+func writeTree(t *testing.T, archive string) string {
+	root := t.TempDir()
+
+	for _, f := range txtar.Parse([]byte(archive)).Files {
+		filename := filepath.Join(root, f.Name)
+		if err := os.MkdirAll(filepath.Dir(filename), 0777); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.WriteFile(filename, f.Data, 0666); err != nil {
+			t.Fatal(err)
+		}
+	}
+	return root
+}

gopls/internal/cache/load.go

@@ -194,7 +194,7 @@ func (s *Snapshot) load(ctx context.Context, allowNetwork bool, scopes ...loadSc
 					return errorf("go/packages returned multiple standalone packages")
 				}
 				standalonePkg = pkg
-			} else if packagesinternal.GetForTest(pkg) == "" && !strings.HasSuffix(pkg.ID, ".test") {
+			} else if pkg.ForTest == "" && !strings.HasSuffix(pkg.ID, ".test") {
 				return errorf("go/packages returned unexpected package %q for standalone file", pkg.ID)
 			}
 		}
@@ -259,7 +259,7 @@ func (s *Snapshot) load(ctx context.Context, allowNetwork bool, scopes ...loadSc
 			s.setBuiltin(pkg.GoFiles[0])
 			continue
 		}
-		if packagesinternal.GetForTest(pkg) == "builtin" {
+		if pkg.ForTest == "builtin" {
 			// We don't care about test variants of builtin. This caused test
 			// failures in https://go.dev/cl/620196, when a test file was added to
 			// builtin.
@@ -416,7 +416,7 @@ func buildMetadata(updates map[PackageID]*metadata.Package, pkg *packages.Packag
 		ID:         id,
 		PkgPath:    pkgPath,
 		Name:       PackageName(pkg.Name),
-		ForTest:    PackagePath(packagesinternal.GetForTest(pkg)),
+		ForTest:    PackagePath(pkg.ForTest),
 		TypesSizes: pkg.TypesSizes,
 		LoadDir:    loadDir,
 		Module:     pkg.Module,

gopls/internal/cache/snapshot.go

@@ -385,7 +385,7 @@ func (s *Snapshot) config(ctx context.Context, inv *gocommand.Invocation) *packa
 			packages.NeedModule |
 			packages.NeedEmbedFiles |
 			packages.LoadMode(packagesinternal.DepsErrors) |
-			packages.LoadMode(packagesinternal.ForTest),
+			packages.NeedForTest,
 		Fset:    nil, // we do our own parsing
 		Overlay: s.buildOverlays(),
 		ParseFile: func(*token.FileSet, string, []byte) (*ast.File, error) {

gopls/internal/cache/typerefs/pkgrefs_test.go

@@ -342,7 +342,7 @@ func loadPackages(query string, needExport bool) (map[PackageID]string, Metadata
 			packages.NeedModule |
 			packages.NeedEmbedFiles |
 			packages.LoadMode(packagesinternal.DepsErrors) |
-			packages.LoadMode(packagesinternal.ForTest),
+			packages.NeedForTest,
 		Tests: true,
 	}
 	if needExport {
@@ -364,7 +364,7 @@ func loadPackages(query string, needExport bool) (map[PackageID]string, Metadata
 			ID:         id,
 			PkgPath:    PackagePath(pkg.PkgPath),
 			Name:       packageName(pkg.Name),
-			ForTest:    PackagePath(packagesinternal.GetForTest(pkg)),
+			ForTest:    PackagePath(pkg.ForTest),
 			TypesSizes: pkg.TypesSizes,
 			LoadDir:    cfg.Dir,
 			Module:     pkg.Module,

internal/drivertest/driver_test.go

@@ -68,7 +68,7 @@ package lib
 			packages.NeedModule |
 			packages.NeedEmbedFiles |
 			packages.LoadMode(packagesinternal.DepsErrors) |
-			packages.LoadMode(packagesinternal.ForTest),
+			packages.NeedForTest,
 	}
 
 	tests := []struct {

internal/packagesinternal/packages.go

@@ -5,7 +5,6 @@
 // Package packagesinternal exposes internal-only fields from go/packages.
 package packagesinternal
 
-var GetForTest = func(p interface{}) string { return "" }
 var GetDepsErrors = func(p interface{}) []*PackageError { return nil }
 
 type PackageError struct {
@@ -16,7 +15,6 @@ type PackageError struct {
 
 var TypecheckCgo int
 var DepsErrors int // must be set as a LoadMode to call GetDepsErrors
-var ForTest int    // must be set as a LoadMode to call GetForTest
 
 var SetModFlag = func(config interface{}, value string) {}
 var SetModFile = func(config interface{}, value string) {}