From d75fbbaadb3239b19063a4d5eef9fc81d0e4345c Mon Sep 17 00:00:00 2001
From: Charlie Vieth <charlie.vieth@gmail.com>
Date: Tue, 24 May 2022 21:46:41 -0400
Subject: [PATCH] buildutil: add MatchFile to replicate build.Context.MatchFile

MatchFile is like build.Context.MatchFile but takes a "src" argument and
returns the package name of the parsed file.
---
 buildutil.go      | 42 ++++++++++++++++++++++--
 buildutil_test.go | 84 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 3 deletions(-)

diff --git a/buildutil.go b/buildutil.go
index d10933c..a75259a 100644
--- a/buildutil.go
+++ b/buildutil.go
@@ -170,7 +170,35 @@ func ReadImports(path string, src interface{}) (pkgname string, imports []string
 	return
 }
 
-func openReader(ctxt *build.Context, filename string, src interface{}) (io.ReadCloser, error) {
+// MatchFile reports whether the file with the given name matches the context
+// and would be included in a Package created by ImportDir. It also returns
+// the package name of the file.
+//
+// MatchFile considers the name of the file and may use ctxt.OpenFile to
+// read some or all of the file's content. If src is not nil it will be
+// used as the content of the file.
+func MatchFile(ctxt *build.Context, dir, name string, src interface{}) (pkgName string, match bool, err error) {
+	rc, err := openReaderDirName(ctxt, dir, name, src)
+	if err != nil {
+		return
+	}
+	data, err := readImportsFast(rc)
+	rc.Close()
+	if err != nil {
+		return "", false, err
+	}
+	pkgName, err = readPackageName(data)
+	if err != nil {
+		return "", false, err
+	}
+	if !GoodOSArchFile(ctxt, name, nil) {
+		return pkgName, false, nil
+	}
+	match, _, err = shouldBuild(ctxt, data, nil)
+	return
+}
+
+func openReaderDirName(ctxt *build.Context, dir, name string, src interface{}) (io.ReadCloser, error) {
 	if src != nil {
 		switch s := src.(type) {
 		case string:
@@ -185,10 +213,18 @@ func openReader(ctxt *build.Context, filename string, src interface{}) (io.ReadC
 			return nil, errors.New("invalid source")
 		}
 	}
+	// If dir is not empty it is joined with name
+	if dir != "" {
+		name = joinPath(ctxt, dir, name)
+	}
 	if ctxt.OpenFile != nil {
-		return ctxt.OpenFile(filename)
+		return ctxt.OpenFile(name)
 	}
-	return os.Open(filename)
+	return os.Open(name)
+}
+
+func openReader(ctxt *build.Context, filename string, src interface{}) (io.ReadCloser, error) {
+	return openReaderDirName(ctxt, "", filename, src)
 }
 
 var (
diff --git a/buildutil_test.go b/buildutil_test.go
index ac21807..0f6ad35 100644
--- a/buildutil_test.go
+++ b/buildutil_test.go
@@ -591,6 +591,71 @@ func TestParseBuildConstraint(t *testing.T) {
 	}
 }
 
+func testMatchFile(t *testing.T, ctxt *build.Context, dir string) {
+	des, err := os.ReadDir(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if testing.Short() && len(des) > 64 {
+		des = des[:64]
+	}
+
+	t.Cleanup(func() {
+		if t.Failed() {
+			t.Logf("Context: %+v", *ctxt)
+		}
+	})
+	for _, d := range des {
+		name := d.Name()
+		if !d.Type().IsRegular() || !strings.HasSuffix(name, ".go") {
+			continue
+		}
+		if testing.Short() && strings.HasSuffix(name, "_test.go") {
+			continue
+		}
+		gotName, gotMatch, err := MatchFile(ctxt, dir, name, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		wantName, err := ReadPackageName(filepath.Join(dir, name), nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		wantMatch, err := ctxt.MatchFile(dir, name)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if gotMatch != wantMatch || (wantMatch && gotName != wantName) {
+			t.Errorf("MatchFile(%q) = %q, %t; want: %q, %t", name,
+				gotName, gotMatch, wantName, wantMatch)
+		}
+	}
+}
+
+func TestMatchFile(t *testing.T) {
+	dir := filepath.Join(runtime.GOROOT(), "src", "runtime")
+	if fi, err := os.Stat(dir); err != nil || !fi.IsDir() {
+		t.Skip("skipping: test requires Go source:", err)
+	}
+
+	t.Run("darwin/arm64", func(t *testing.T) {
+		t.Parallel()
+		ctxt := build.Default
+		ctxt.GOOS = "darwin"
+		ctxt.GOARCH = "arm64"
+		testMatchFile(t, &ctxt, dir)
+	})
+
+	t.Run("linux/amd64", func(t *testing.T) {
+		t.Parallel()
+		ctxt := build.Default
+		ctxt.GOOS = "linux"
+		ctxt.GOARCH = "amd64"
+		ctxt.CgoEnabled = true
+		testMatchFile(t, &ctxt, dir)
+	})
+}
+
 func BenchmarkImportPath(b *testing.B) {
 	wd, err := os.Getwd()
 	if err != nil {
@@ -686,3 +751,22 @@ func BenchmarkShortImport_Overlay(b *testing.B) {
 
 	benchmarkShortImport(b, &ctxt, list)
 }
+
+func BenchmarkMatchFile(b *testing.B) {
+	dir := b.TempDir()
+	name := filepath.Join(dir, "build.go")
+	// if err := os.WriteFile(name, []byte(LongPackageHeader), 0644); err != nil {
+	const content = "package foo\n"
+	if err := os.WriteFile(name, []byte(content), 0644); err != nil {
+		b.Fatal(err)
+	}
+	ctxt := build.Default
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, _, err := MatchFile(&ctxt, dir, name, LongPackageHeader)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}