From 55a956e1897b7ac5aa247760f955ff208b8946f5 Mon Sep 17 00:00:00 2001
From: Ricky Stewart <ricky@cockroachlabs.com>
Date: Fri, 23 Feb 2024 16:02:12 -0600
Subject: [PATCH] gcassert: allow absolute paths to be used for packages, and
 add cwd arg

---
 gcassert.go      | 26 +++++++++++++++++++-----
 gcassert_test.go | 53 ++++++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/gcassert.go b/gcassert.go
index bd10056..51254f5 100644
--- a/gcassert.go
+++ b/gcassert.go
@@ -144,8 +144,24 @@ func (v assertVisitor) Visit(node ast.Node) (w ast.Visitor) {
 // GCAssert searches through the packages at the input path and writes failures
 // to comply with //gcassert directives to the given io.Writer.
 func GCAssert(w io.Writer, paths ...string) error {
+	return GCAssertCwd(w, "", paths...)
+}
+
+// GCAssertCwd performs the same operation as GCAssert, but runs `go build` in
+// the provided working directory `cwd`. If `cwd` is the empty string, then
+/// `go build` will be run in the current working directory.
+func GCAssertCwd(w io.Writer, cwd string, paths ...string) error {
+	var err error
+	if cwd == "" {
+		cwd, err = os.Getwd()
+		if err != nil {
+			return err
+		}
+	}
+
 	fileSet := token.NewFileSet()
 	pkgs, err := packages.Load(&packages.Config{
+		Dir: cwd,
 		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedCompiledGoFiles |
 			packages.NeedTypesInfo | packages.NeedTypes,
 		Fset: fileSet,
@@ -160,13 +176,13 @@ func GCAssert(w io.Writer, paths ...string) error {
 
 	args := []string{"build", "-gcflags=-m=2 -d=ssa/check_bce/debug=1"}
 	for i := range paths {
-		args = append(args, "./"+paths[i])
+		if filepath.IsAbs(paths[i]) {
+			args = append(args, paths[i])
+		} else {
+			args = append(args, "./"+paths[i])
+		}
 	}
 	cmd := exec.Command("go", args...)
-	cwd, err := os.Getwd()
-	if err != nil {
-		return err
-	}
 	cmd.Dir = cwd
 	pr, pw := io.Pipe()
 	// Create a temp file to log all diagnostic output.
diff --git a/gcassert_test.go b/gcassert_test.go
index 79f8a87..fd43103 100644
--- a/gcassert_test.go
+++ b/gcassert_test.go
@@ -77,12 +77,10 @@ func TestParseDirectives(t *testing.T) {
 }
 
 func TestGCAssert(t *testing.T) {
-	var w strings.Builder
-	err := GCAssert(&w, "./testdata", "./testdata/otherpkg")
+	cwd, err := os.Getwd()
 	if err != nil {
 		t.Fatal(err)
 	}
-
 	expectedOutput := `testdata/noescape.go:21:	foo := foo{a: 1, b: 2}: foo escapes to heap:
 testdata/noescape.go:35:	// This annotation should fail, because f will escape to the heap.
 //
@@ -103,5 +101,52 @@ testdata/inline.go:58:	test(0).neverInlinedMethod(10): call was not inlined
 testdata/inline.go:60:	otherpkg.A{}.NeverInlined(sum): call was not inlined
 testdata/issue5.go:4:	Gen().Layout(): call was not inlined
 `
-	assert.Equal(t, expectedOutput, w.String())
+
+	testCases := []struct{
+		name string
+		pkgs []string
+		cwd string
+		expected string
+	}{
+		{
+			name: "relative",
+			pkgs: []string{
+				"./testdata",
+				"./testdata/otherpkg",
+			},
+			expected: expectedOutput,
+		},
+		{
+			name: "absolute",
+			pkgs: []string{
+				filepath.Join(cwd, "testdata"),
+				filepath.Join(cwd, "testdata/otherpkg"),
+			},
+			expected: expectedOutput,
+		},
+		{
+			name: "relative-cwd",
+			pkgs: []string{
+				".",
+				"./otherpkg",
+			},
+			cwd: filepath.Join(cwd, "testdata"),
+			expected: strings.ReplaceAll(expectedOutput, "testdata/", ""),
+		},
+	}
+	for _, testCase := range testCases {
+		var w strings.Builder
+		t.Run(testCase.name, func(t *testing.T) {
+			var err error
+			if testCase.cwd == "" {
+				err = GCAssert(&w, testCase.pkgs...)
+			} else {
+				err = GCAssertCwd(&w, testCase.cwd, testCase.pkgs...)
+			}
+			if err != nil {
+				t.Fatal(err)
+			}
+			assert.Equal(t, testCase.expected, w.String())
+		})
+	}
 }