From b6647bb2159ca38846b38e92c5b7a62dd66d6a13 Mon Sep 17 00:00:00 2001
From: healthy-pod <ahmad@cockroachlabs.com>
Date: Mon, 15 May 2023 18:27:16 -0700
Subject: [PATCH] Add support for building using `bazel`

This code change adds support for building using `bazel`
if the `--use-bazel` flag is set. It will be unset by
default.
---
 README.md            |  2 +-
 cmd/gcassert/main.go |  4 +++-
 gcassert.go          | 35 ++++++++++++++++++++++++++++-------
 gcassert_test.go     |  2 +-
 4 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 30f90df..9c52933 100644
--- a/README.md
+++ b/README.md
@@ -109,7 +109,7 @@ import "github.com/jordanlewis/gcassert"
 
 func main() {
     var buf strings.Builder
-    if err := gcassert.GCAssert(&buf, "./path/to/package", "./otherpath/to/package"); err != nil {
+    if err := gcassert.GCAssert(&buf, false /* useBazel */ "./path/to/package", "./otherpath/to/package"); err != nil {
         // handle non-lint-failure related errors
         panic(err)
     }
diff --git a/cmd/gcassert/main.go b/cmd/gcassert/main.go
index 136f38e..62a155b 100644
--- a/cmd/gcassert/main.go
+++ b/cmd/gcassert/main.go
@@ -9,10 +9,12 @@ import (
 	"github.com/jordanlewis/gcassert"
 )
 
+var useBazel = flag.Bool("use-bazel", false, "use bazel to build")
+
 func main() {
 	flag.Parse()
 	var buf strings.Builder
-	err := gcassert.GCAssert(&buf, flag.Args()...)
+	err := gcassert.GCAssert(&buf, *useBazel, flag.Args()...)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
diff --git a/gcassert.go b/gcassert.go
index f0ed461..0a23e97 100644
--- a/gcassert.go
+++ b/gcassert.go
@@ -2,7 +2,6 @@ package gcassert
 
 import (
 	"bufio"
-	"errors"
 	"fmt"
 	"go/ast"
 	"go/printer"
@@ -38,7 +37,7 @@ func stringToDirective(s string) (assertDirective, error) {
 	case "noescape":
 		return noescape, nil
 	}
-	return noDirective, errors.New(fmt.Sprintf("no such directive %s", s))
+	return noDirective, fmt.Errorf("no such directive %s", s)
 }
 
 // passInfo contains info on a passed directive for directives that have
@@ -143,13 +142,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 {
+func GCAssert(w io.Writer, useBazel bool, paths ...string) error {
+	for _, path := range paths {
+		// Assert that all paths begin with './'
+		// This is needed for ensuring that packages.Load and parseDirectives work
+		// as expected.
+		if !strings.HasPrefix(path, "./") {
+			return fmt.Errorf("all paths should be prefixed with './': got %s", path)
+		}
+	}
 	fileSet := token.NewFileSet()
 	pkgs, err := packages.Load(&packages.Config{
 		Mode: packages.NeedName | packages.NeedFiles | packages.NeedSyntax | packages.NeedCompiledGoFiles |
 			packages.NeedTypesInfo | packages.NeedTypes,
 		Fset: fileSet,
 	}, paths...)
+	if err != nil {
+		return err
+	}
 	directiveMap, err := parseDirectives(pkgs, fileSet)
 	if err != nil {
 		return err
@@ -158,11 +168,22 @@ func GCAssert(w io.Writer, paths ...string) error {
 	// Next: invoke Go compiler with -m flags to get the compiler to print
 	// its optimization decisions.
 
-	args := []string{"build", "-gcflags=-m=2 -d=ssa/check_bce/debug=1"}
-	for i := range paths {
-		args = append(args, "./"+paths[i])
+	var args []string
+	var cmd *exec.Cmd
+
+	if useBazel {
+		args = []string{"build"}
+		for i := range paths {
+			args = append(args, strings.TrimPrefix(paths[i], "./"))
+		}
+		args = append(args, "--@io_bazel_rules_go//go/config:gc_goopts=-m=2,-d=ssa/check_bce/debug=1")
+		cmd = exec.Command("bazel", args...)
+	} else {
+		args = []string{"build", "-gcflags=-m=2 -d=ssa/check_bce/debug=1"}
+		args = append(args, paths...)
+		cmd = exec.Command("go", args...)
 	}
-	cmd := exec.Command("go", args...)
+
 	cwd, err := os.Getwd()
 	if err != nil {
 		return err
diff --git a/gcassert_test.go b/gcassert_test.go
index 7fa58c7..22b68f0 100644
--- a/gcassert_test.go
+++ b/gcassert_test.go
@@ -77,7 +77,7 @@ func TestParseDirectives(t *testing.T) {
 
 func TestGCAssert(t *testing.T) {
 	var w strings.Builder
-	err := GCAssert(&w, "./testdata", "./testdata/otherpkg")
+	err := GCAssert(&w, false /* useBazel */, "./testdata", "./testdata/otherpkg")
 	if err != nil {
 		t.Fatal(err)
 	}