Refactor nogo to better support ARM.

PiperOrigin-RevId: 337544107

4 files changed, 109 insertions, 22 deletions

diff --git a/tools/bazeldefs/defs.bzl b/tools/bazeldefs/defs.bzl
index cf5b1dc0d..4aa55e8e5 100644
--- a/tools/bazeldefs/defs.bzl
+++ b/tools/bazeldefs/defs.bzl
@@ -158,11 +158,31 @@ def go_test_library(target):
         return target.attr.embed[0]
     return None
 
-def go_context(ctx, std = False):
+def go_context(ctx, goos = None, goarch = None, std = False):
+    """Extracts a standard Go context struct.
+
+    Args:
+      ctx: the starlark context (required).
+      goos: the GOOS value.
+      goarch: the GOARCH value.
+      std: ignored.
+
+    Returns:
+      A context Go struct with pointers to Go toolchain components.
+    """
+
     # We don't change anything for the standard library analysis. All Go files
     # are available in all instances. Note that this includes the standard
     # library sources, which are analyzed by nogo.
     go_ctx = _go_context(ctx)
+    if goos == None:
+        goos = go_ctx.sdk.goos
+    elif goos != go_ctx.sdk.goos:
+        fail("Internal GOOS (%s) doesn't match GoSdk GOOS (%s)." % (goos, go_ctx.sdk.goos))
+    if goarch == None:
+        goarch = go_ctx.sdk.goarch
+    elif goarch != go_ctx.sdk.goarch:
+        fail("Internal GOARCH (%s) doesn't match GoSdk GOARCH (%s)." % (goarch, go_ctx.sdk.goarch))
     return struct(
         go = go_ctx.go,
         env = go_ctx.env,
@@ -186,6 +206,12 @@ def select_arch(amd64 = "amd64", arm64 = "arm64", default = None, **kwargs):
 def select_system(linux = ["__linux__"], **kwargs):
     return linux  # Only Linux supported.
 
+def select_goarch():
+    return select_arch(arm64 = "arm64", amd64 = "amd64")
+
+def select_goos():
+    return select_system(linux = "linux")
+
 def default_installer():
     return None
 
diff --git a/tools/defs.bzl b/tools/defs.bzl
index e2c72a1f6..9a90f2554 100644
--- a/tools/defs.bzl
+++ b/tools/defs.bzl
@@ -7,7 +7,7 @@ change for Google-internal and bazel-compatible rules.
 
 load("//tools/go_stateify:defs.bzl", "go_stateify")
 load("//tools/go_marshal:defs.bzl", "go_marshal", "marshal_deps", "marshal_test_deps")
-load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_grpc_library = "cc_grpc_library", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _gazelle = "gazelle", _gbenchmark = "gbenchmark", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_test = "go_test", _grpcpp = "grpcpp", _gtest = "gtest", _loopback = "loopback", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar", _proto_library = "proto_library", _py_binary = "py_binary", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path", _vdso_linker_option = "vdso_linker_option")
+load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_grpc_library = "cc_grpc_library", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _gazelle = "gazelle", _gbenchmark = "gbenchmark", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_test = "go_test", _grpcpp = "grpcpp", _gtest = "gtest", _loopback = "loopback", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar", _proto_library = "proto_library", _py_binary = "py_binary", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_goarch = "select_goarch", _select_goos = "select_goos", _select_system = "select_system", _short_path = "short_path", _vdso_linker_option = "vdso_linker_option")
 load("//tools/bazeldefs:platforms.bzl", _default_platform = "default_platform", _platforms = "platforms")
 load("//tools/bazeldefs:tags.bzl", "go_suffixes")
 load("//tools/nogo:defs.bzl", "nogo_test")
@@ -35,6 +35,8 @@ pkg_tar = _pkg_tar
 py_binary = _py_binary
 select_arch = _select_arch
 select_system = _select_system
+select_goos = _select_goos
+select_goarch = _select_goarch
 short_path = _short_path
 rbe_platform = _rbe_platform
 rbe_toolchain = _rbe_toolchain
diff --git a/tools/nogo/BUILD b/tools/nogo/BUILD
index dd4b46f58..3c6be3339 100644
--- a/tools/nogo/BUILD
+++ b/tools/nogo/BUILD
@@ -1,8 +1,15 @@
-load("//tools:defs.bzl", "bzl_library", "go_library")
-load("//tools/nogo:defs.bzl", "nogo_objdump_tool", "nogo_stdlib")
+load("//tools:defs.bzl", "bzl_library", "go_library", "select_goarch", "select_goos")
+load("//tools/nogo:defs.bzl", "nogo_objdump_tool", "nogo_stdlib", "nogo_target")
 
 package(licenses = ["notice"])
 
+nogo_target(
+    name = "target",
+    goarch = select_goarch(),
+    goos = select_goos(),
+    visibility = ["//visibility:public"],
+)
+
 nogo_objdump_tool(
     name = "objdump_tool",
     visibility = ["//visibility:public"],
diff --git a/tools/nogo/defs.bzl b/tools/nogo/defs.bzl
index c6fcfd402..8096087dd 100644
--- a/tools/nogo/defs.bzl
+++ b/tools/nogo/defs.bzl
@@ -2,9 +2,33 @@
 
 load("//tools/bazeldefs:defs.bzl", "go_context", "go_importpath", "go_rule", "go_test_library")
 
-def _nogo_objdump_tool_impl(ctx):
-    go_ctx = go_context(ctx)
+NogoTargetInfo = provider(
+    "information about the Go target",
+    fields = {
+        "goarch": "the build architecture (GOARCH)",
+        "goos": "the build OS target (GOOS)",
+    },
+)
+
+def _nogo_target_impl(ctx):
+    return [NogoTargetInfo(
+        goarch = ctx.attr.goarch,
+        goos = ctx.attr.goos,
+    )]
 
+nogo_target = go_rule(
+    rule,
+    implementation = _nogo_target_impl,
+    attrs = {
+        # goarch is the build architecture. This will normally be provided by a
+        # select statement, but this information is propagated to other rules.
+        "goarch": attr.string(mandatory = True),
+        # goos is similarly the build operating system target.
+        "goos": attr.string(mandatory = True),
+    },
+)
+
+def _nogo_objdump_tool_impl(ctx):
     # Construct the magic dump command.
     #
     # Note that in some cases, the input is being fed into the tool via stdin.
@@ -12,6 +36,8 @@ def _nogo_objdump_tool_impl(ctx):
     # we need the tool to handle this case by creating a temporary file.
     #
     # [1] https://github.com/golang/go/issues/41051
+    nogo_target_info = ctx.attr._nogo_target[NogoTargetInfo]
+    go_ctx = go_context(ctx, goos = nogo_target_info.goos, goarch = nogo_target_info.goarch)
     env_prefix = " ".join(["%s=%s" % (key, value) for (key, value) in go_ctx.env.items()])
     dumper = ctx.actions.declare_file(ctx.label.name)
     ctx.actions.write(dumper, "\n".join([
@@ -42,6 +68,12 @@ def _nogo_objdump_tool_impl(ctx):
 nogo_objdump_tool = go_rule(
     rule,
     implementation = _nogo_objdump_tool_impl,
+    attrs = {
+        "_nogo_target": attr.label(
+            default = "//tools/nogo:target",
+            cfg = "target",
+        ),
+    },
 )
 
 # NogoStdlibInfo is the set of standard library facts.
@@ -54,9 +86,9 @@ NogoStdlibInfo = provider(
 )
 
 def _nogo_stdlib_impl(ctx):
-    go_ctx = go_context(ctx)
-
     # Build the standard library facts.
+    nogo_target_info = ctx.attr._nogo_target[NogoTargetInfo]
+    go_ctx = go_context(ctx, goos = nogo_target_info.goos, goarch = nogo_target_info.goarch)
     facts = ctx.actions.declare_file(ctx.label.name + ".facts")
     findings = ctx.actions.declare_file(ctx.label.name + ".findings")
     config = struct(
@@ -70,12 +102,12 @@ def _nogo_stdlib_impl(ctx):
     ctx.actions.run(
         inputs = [config_file] + go_ctx.stdlib_srcs,
         outputs = [facts, findings],
-        tools = depset(go_ctx.runfiles.to_list() + ctx.files._objdump_tool),
-        executable = ctx.files._nogo[0],
+        tools = depset(go_ctx.runfiles.to_list() + ctx.files._nogo_objdump_tool),
+        executable = ctx.files._nogo_check[0],
         mnemonic = "GoStandardLibraryAnalysis",
         progress_message = "Analyzing Go Standard Library",
         arguments = go_ctx.nogo_args + [
-            "-objdump_tool=%s" % ctx.files._objdump_tool[0].path,
+            "-objdump_tool=%s" % ctx.files._nogo_objdump_tool[0].path,
             "-stdlib=%s" % config_file.path,
             "-findings=%s" % findings.path,
             "-facts=%s" % facts.path,
@@ -92,11 +124,17 @@ nogo_stdlib = go_rule(
     rule,
     implementation = _nogo_stdlib_impl,
     attrs = {
-        "_nogo": attr.label(
+        "_nogo_check": attr.label(
             default = "//tools/nogo/check:check",
+            cfg = "host",
         ),
-        "_objdump_tool": attr.label(
+        "_nogo_objdump_tool": attr.label(
             default = "//tools/nogo:objdump_tool",
+            cfg = "host",
+        ),
+        "_nogo_target": attr.label(
+            default = "//tools/nogo:target",
+            cfg = "target",
         ),
     },
 )
@@ -119,8 +157,6 @@ NogoInfo = provider(
 )
 
 def _nogo_aspect_impl(target, ctx):
-    go_ctx = go_context(ctx)
-
     # If this is a nogo rule itself (and not the shadow of a go_library or
     # go_binary rule created by such a rule), then we simply return nothing.
     # All work is done in the shadow properties for go rules. For a proto
@@ -200,10 +236,13 @@ def _nogo_aspect_impl(target, ctx):
         inputs += info.binaries
 
     # Add the standard library facts.
-    stdlib_facts = ctx.attr._nogo_stdlib[NogoStdlibInfo].facts
+    stdlib_info = ctx.attr._nogo_stdlib[NogoStdlibInfo]
+    stdlib_facts = stdlib_info.facts
     inputs.append(stdlib_facts)
 
     # The nogo tool operates on a configuration serialized in JSON format.
+    nogo_target_info = ctx.attr._nogo_target[NogoTargetInfo]
+    go_ctx = go_context(ctx, goos = nogo_target_info.goos, goarch = nogo_target_info.goarch)
     facts = ctx.actions.declare_file(target.label.name + ".facts")
     findings = ctx.actions.declare_file(target.label.name + ".findings")
     escapes = ctx.actions.declare_file(target.label.name + ".escapes")
@@ -224,13 +263,13 @@ def _nogo_aspect_impl(target, ctx):
     ctx.actions.run(
         inputs = inputs,
         outputs = [facts, findings, escapes],
-        tools = depset(go_ctx.runfiles.to_list() + ctx.files._objdump_tool),
-        executable = ctx.files._nogo[0],
+        tools = depset(go_ctx.runfiles.to_list() + ctx.files._nogo_objdump_tool),
+        executable = ctx.files._nogo_check[0],
         mnemonic = "GoStaticAnalysis",
         progress_message = "Analyzing %s" % target.label,
         arguments = go_ctx.nogo_args + [
             "-binary=%s" % target_objfile.path,
-            "-objdump_tool=%s" % ctx.files._objdump_tool[0].path,
+            "-objdump_tool=%s" % ctx.files._nogo_objdump_tool[0].path,
             "-package=%s" % config_file.path,
             "-findings=%s" % findings.path,
             "-facts=%s" % facts.path,
@@ -266,9 +305,22 @@ nogo_aspect = go_rule(
         "embed",
     ],
     attrs = {
-        "_nogo": attr.label(default = "//tools/nogo/check:check"),
-        "_nogo_stdlib": attr.label(default = "//tools/nogo:stdlib"),
-        "_objdump_tool": attr.label(default = "//tools/nogo:objdump_tool"),
+        "_nogo_check": attr.label(
+            default = "//tools/nogo/check:check",
+            cfg = "host",
+        ),
+        "_nogo_stdlib": attr.label(
+            default = "//tools/nogo:stdlib",
+            cfg = "host",
+        ),
+        "_nogo_objdump_tool": attr.label(
+            default = "//tools/nogo:objdump_tool",
+            cfg = "host",
+        ),
+        "_nogo_target": attr.label(
+            default = "//tools/nogo:target",
+            cfg = "target",
+        ),
     },
 )