Apply map_kind to args as well as rule kinds

bazel-contrib#1310 added support for adding args from Gazelle, so you can express things like: ``` maybe( java_test, ... ) ``` But `map_kind` doesn't currently get applied to these args. This PR adds that application.
illicitonion · Apr 2, 2024 · 17c1110 · 17c1110
1 parent 34bea0a
commit 17c1110
Show file tree

Hide file tree

Showing 5 changed files with 295 additions and 6 deletions.
diff --git a/cmd/gazelle/BUILD.bazel b/cmd/gazelle/BUILD.bazel
@@ -36,6 +36,7 @@ go_library(
         "//resolve",
         "//rule",
         "//walk",
+        "@com_github_bazelbuild_buildtools//build",
         "@com_github_pmezard_go_difflib//difflib",
     ],
 )

diff --git a/cmd/gazelle/fix-update.go b/cmd/gazelle/fix-update.go
@@ -38,6 +38,7 @@ import (
 	"github.com/bazelbuild/bazel-gazelle/resolve"
 	"github.com/bazelbuild/bazel-gazelle/rule"
 	"github.com/bazelbuild/bazel-gazelle/walk"
+	"github.com/bazelbuild/buildtools/build"
 )
 
 // updateConfig holds configuration information needed to run the fix and
@@ -369,17 +370,45 @@ func runFixUpdate(wd string, cmd command, args []string) (err error) {
 		if f != nil {
 			allRules = append(allRules, f.Rules...)
 		}
-		for _, r := range allRules {
-			repl, err := lookupMapKindReplacement(c.KindMap, r.Kind())
+
+		maybeRecordReplacement := func(ruleKind string) (*string, error) {
+			repl, err := lookupMapKindReplacement(c.KindMap, ruleKind)
 			if err != nil {
-				errorsFromWalk = append(errorsFromWalk, fmt.Errorf("looking up mapped kind: %w", err))
-				continue
+				return nil, err
 			}
 			if repl != nil {
-				mappedKindInfo[repl.KindName] = kinds[r.Kind()]
+				mappedKindInfo[repl.KindName] = kinds[ruleKind]
 				mappedKinds = append(mappedKinds, *repl)
 				mrslv.MappedKind(rel, *repl)
-				r.SetKind(repl.KindName)
+				return &repl.KindName, nil
+			}
+			return nil, nil
+		}
+
+		for _, r := range allRules {
+			if replacementName, err := maybeRecordReplacement(r.Kind()); err != nil {
+				errorsFromWalk = append(errorsFromWalk, fmt.Errorf("looking up mapped kind: %w", err))
+			} else if replacementName != nil {
+				r.SetKind(*replacementName)
+			}
+
+			for i, arg := range r.Args() {
+				// Only check the first arg - this supports the maybe(java_library, ...) pattern,
+				// but avoids potential false positives from other uses of symbols.
+				if i != 0 {
+					break
+				}
+				if ident, ok := arg.(*build.Ident); ok {
+					// Don't allow re-mapping symbols that aren't known loads of a plugin.
+					if _, knownKind := kinds[ident.Name]; !knownKind {
+						continue
+					}
+					if replacementName, err := maybeRecordReplacement(ident.Name); err != nil {
+						errorsFromWalk = append(errorsFromWalk, fmt.Errorf("looking up mapped kind: %w", err))
+					} else if replacementName != nil {
+						r.UpdateArg(i, &build.Ident{Name: *replacementName})
+					}
+				}
 			}
 		}
 		for _, r := range empty {

diff --git a/cmd/gazelle/fix_test.go b/cmd/gazelle/fix_test.go
@@ -407,3 +407,250 @@ go_library(
 
 	testtools.CheckFiles(t, dir, fixture)
 }
+
+func TestFix_MapKind_Argument(t *testing.T) {
+	for name, tc := range map[string]struct {
+		before []testtools.FileSpec
+		after  []testtools.FileSpec
+	}{
+		"same-name": {
+			before: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+
+# gazelle:map_kind go_binary go_binary //my:custom.bzl
+
+maybe(
+    go_binary,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)`,
+				},
+				{
+					Path:    "some.go",
+					Content: `package some`,
+				},
+			},
+			after: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("//my:custom.bzl", "go_binary")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+# gazelle:map_kind go_binary go_binary //my:custom.bzl
+
+maybe(
+    go_binary,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)`,
+				},
+			},
+		},
+		"different-name": {
+			before: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+
+# gazelle:map_kind go_binary custom_go_binary //my:custom.bzl
+
+maybe(
+    go_binary,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)`,
+				},
+				{
+					Path:    "some.go",
+					Content: `package some`,
+				},
+			},
+			after: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("//my:custom.bzl", "custom_go_binary")
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+# gazelle:map_kind go_binary custom_go_binary //my:custom.bzl
+
+maybe(
+    custom_go_binary,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)`,
+				},
+			},
+		},
+		"non-loaded-symbol": {
+			before: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+custom_go_library = go_library
+
+# This will be ignored because it's not a directly loaded symbol when used:
+# gazelle:map_kind custom_go_library custom_go_library //my:custom.bzl
+
+maybe(
+    custom_go_library,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)
+`,
+				},
+				{
+					Path:    "some.go",
+					Content: `package some`,
+				},
+			},
+			after: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+
+custom_go_library = go_library
+
+# This will be ignored because it's not a directly loaded symbol when used:
+# gazelle:map_kind custom_go_library custom_go_library //my:custom.bzl
+
+maybe(
+    custom_go_library,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)
+`,
+				},
+			},
+		},
+		"not-arg-0": {
+			before: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@other_rules//:def.bzl", "something_custom")
+
+# gazelle:map_kind something_custom something_custom //my:custom.bzl
+
+maybe(
+    go_library,
+    something_custom,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)
+`,
+				},
+				{
+					Path:    "some.go",
+					Content: `package some`,
+				},
+			},
+			after: []testtools.FileSpec{
+				{
+					Path: "BUILD.bazel",
+					Content: `
+load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@other_rules//:def.bzl", "something_custom")
+
+# gazelle:map_kind something_custom something_custom //my:custom.bzl
+
+maybe(
+    go_library,
+    something_custom,
+    name = "nofix",
+    library = ":go_default_library",
+    visibility = ["//visibility:public"],
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = ["some.go"],
+    importpath = "example.com/repo",
+    visibility = ["//visibility:public"],
+)
+`,
+				},
+			},
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			dir, cleanup := testtools.CreateFiles(t, tc.before)
+			defer cleanup()
+
+			if err := run(dir, []string{
+				"-repo_root", dir,
+				"-go_prefix", "example.com/repo",
+				dir,
+			}); err != nil {
+				t.Fatalf("run failed: %v", err)
+			}
+
+			testtools.CheckFiles(t, dir, tc.after)
+		})
+	}
+}
diff --git a/rule/rule.go b/rule/rule.go
@@ -1016,6 +1016,15 @@ func (r *Rule) AddArg(value bzl.Expr) {
 	r.updated = true
 }
 
+// UpdateArg replaces an existing arg with a new value.
+func (r *Rule) UpdateArg(index int, value bzl.Expr) {
+	if len(r.args) < index {
+		panic(fmt.Sprintf("Can't update argument at index %d, only have %d args", index, len(r.args)))
+	}
+	r.args[index] = value
+	r.updated = true
+}
+
 // SortedAttrs returns the keys of attributes whose values will be sorted
 func (r *Rule) SortedAttrs() []string {
 	return r.sortedAttrs

diff --git a/rule/rule_test.go b/rule/rule_test.go
@@ -63,6 +63,8 @@ y_library(name = "bar")
 	loadMaybe.Insert(f, 0)
 	baz := NewRule("maybe", "baz")
 	baz.AddArg(&bzl.LiteralExpr{Token: "z"})
+	baz.AddArg(&bzl.LiteralExpr{Token: "z"})
+	baz.UpdateArg(0, &bzl.LiteralExpr{Token: "z0"})
 	baz.SetAttr("srcs", GlobValue{
 		Patterns: []string{"**"},
 		Excludes: []string{"*.pem"},
@@ -85,6 +87,7 @@ y_library(
 )
 
 maybe(
+    z0,
     z,
     name = "baz",
     srcs = glob(