Support loading rules from rules_go repository

go/tools/gazelle/gazelle/main.go

@@ -28,15 +28,17 @@ import (
 )
 
 var (
-	goPrefix = flag.String("go_prefix", "", "go_prefix of the target workspace")
-	repoRoot = flag.String("repo_root", "", "path to a directory which corresponds to go_prefix")
+	rulesGoRepo = flag.String("rules_go_repo", generator.DefaultRulesGoRepo, "repository name of rules_go")
+	goPrefix    = flag.String("go_prefix", "", "go_prefix of the target workspace")
+	repoRoot    = flag.String("repo_root", "", "path to a directory which corresponds to go_prefix")
 )
 
 func run(dirs []string) error {
 	g, err := generator.New(*repoRoot, *goPrefix)
 	if err != nil {
 		return err
 	}
+	g.RulesGoRepo = *rulesGoRepo
 
 	for _, d := range dirs {
 		files, err := g.Generate(d)

go/tools/gazelle/generator/generator.go

@@ -28,11 +28,21 @@ import (
 	"github.com/bazelbuild/rules_go/go/tools/gazelle/rules"
 )
 
+const (
+	// DefaultRulesGoRepo is the canonical workspace name of rules_go
+	DefaultRulesGoRepo = "@io_bazel_rules_go"
+)
+
 // Generator generates BUILD files for a Go repository.
 type Generator struct {
 	repoRoot string
 	bctx     build.Context
 	g        rules.Generator
+
+	// RulesGoRepo overrides Bazel repository name of rules_go.
+	// It is configurable only for importing external repositories into
+	// rules_go itself. So you usually don't have to specifiy this value.
+	RulesGoRepo string
 }
 
 // New returns a new Generator which is responsible for a Go repository.
@@ -51,9 +61,10 @@ func New(repoRoot, goPrefix string) (*Generator, error) {
 		return nil, err
 	}
 	return &Generator{
-		repoRoot: filepath.Clean(repoRoot),
-		bctx:     bctx,
-		g:        rules.NewGenerator(goPrefix),
+		repoRoot:    filepath.Clean(repoRoot),
+		bctx:        bctx,
+		g:           rules.NewGenerator(goPrefix),
+		RulesGoRepo: DefaultRulesGoRepo,
 	}, nil
 }
 
@@ -81,14 +92,11 @@ func (g *Generator) Generate(dir string) ([]*bzl.File, error) {
 			rel = ""
 		}
 
-		rs, err := g.g.Generate(filepath.ToSlash(rel), pkg)
+		file, err := g.generateOne(rel, pkg)
 		if err != nil {
 			return err
 		}
-		file := &bzl.File{Path: filepath.Join(rel, "BUILD")}
-		for _, r := range rs {
-			file.Stmt = append(file.Stmt, r.Call)
-		}
+
 		files = append(files, file)
 		return nil
 	})
@@ -98,6 +106,54 @@ func (g *Generator) Generate(dir string) ([]*bzl.File, error) {
 	return files, nil
 }
 
+func (g *Generator) generateOne(rel string, pkg *build.Package) (*bzl.File, error) {
+	rs, err := g.g.Generate(filepath.ToSlash(rel), pkg)
+	if err != nil {
+		return nil, err
+	}
+
+	file := &bzl.File{Path: filepath.Join(rel, "BUILD")}
+	for _, r := range rs {
+		file.Stmt = append(file.Stmt, r.Call)
+	}
+	if load := g.generateLoad(file); load != nil {
+		file.Stmt = append([]bzl.Expr{load}, file.Stmt...)
+	}
+	return file, nil
+}
+
+func (g *Generator) generateLoad(f *bzl.File) bzl.Expr {
+	var list []string
+	for _, kind := range []string{
+		"go_prefix",
+		"go_library",
+		"go_binary",
+		"go_test",
+		// TODO(yugui): Support cgo_library
+	} {
+		if len(f.Rules(kind)) > 0 {
+			list = append(list, kind)
+		}
+	}
+	if len(list) == 0 {
+		return nil
+	}
+	return loadExpr(fmt.Sprintf("%s//go:def.bzl", g.RulesGoRepo), list...)
+}
+
+func loadExpr(ruleFile string, rules ...string) bzl.Expr {
+	var list []bzl.Expr
+	for _, r := range append([]string{ruleFile}, rules...) {
+		list = append(list, &bzl.StringExpr{Value: r})
+	}
+
+	return &bzl.CallExpr{
+		X:            &bzl.LiteralExpr{Token: "load"},
+		List:         list,
+		ForceCompact: true,
+	}
+}
+
 func isDescendingDir(dir, root string) bool {
 	if dir == root {
 		return true

go/tools/gazelle/generator/generator_test.go

@@ -32,6 +32,11 @@ func TestGenerator(t *testing.T) {
 	stub := stubRuleGen{
 		fixtures: map[string][]*bzl.Rule{
 			"lib": {
+				{
+					Call: &bzl.CallExpr{
+						X: &bzl.LiteralExpr{Token: "go_prefix"},
+					},
+				},
 				{
 					Call: &bzl.CallExpr{
 						X: &bzl.LiteralExpr{Token: "go_library"},
@@ -44,6 +49,11 @@ func TestGenerator(t *testing.T) {
 						X: &bzl.LiteralExpr{Token: "go_library"},
 					},
 				},
+				{
+					Call: &bzl.CallExpr{
+						X: &bzl.LiteralExpr{Token: "go_test"},
+					},
+				},
 			},
 			"bin": {
 				{
@@ -72,15 +82,26 @@ func TestGenerator(t *testing.T) {
 	want := []*bzl.File{
 		{
 			Path: "lib/BUILD",
-			Stmt: []bzl.Expr{stub.fixtures["lib"][0].Call},
+			Stmt: []bzl.Expr{
+				loadExpr("@io_bazel_rules_go//go:def.bzl", "go_prefix", "go_library"),
+				stub.fixtures["lib"][0].Call,
+				stub.fixtures["lib"][1].Call,
+			},
 		},
 		{
 			Path: "lib/deep/BUILD",
-			Stmt: []bzl.Expr{stub.fixtures["lib/deep"][0].Call},
+			Stmt: []bzl.Expr{
+				loadExpr("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test"),
+				stub.fixtures["lib/deep"][0].Call,
+				stub.fixtures["lib/deep"][1].Call,
+			},
 		},
 		{
 			Path: "bin/BUILD",
-			Stmt: []bzl.Expr{stub.fixtures["bin"][0].Call},
+			Stmt: []bzl.Expr{
+				loadExpr("@io_bazel_rules_go//go:def.bzl", "go_binary"),
+				stub.fixtures["bin"][0].Call,
+			},
 		},
 	}
 	sort.Sort(fileSlice(want))