Automatically propagate tags for stateify and marshal.

Note that files will need to be appropriately segmented in order for the
mechanism to work, in suffixes implying special tags. This only needs to happen
for cases where marshal or state structures are defined, which should be rare
and mostly architecture specific.

PiperOrigin-RevId: 293231579

10 files changed, 235 insertions, 153 deletions

diff --git a/tools/build/defs.bzl b/tools/build/defs.bzl
index 967c1f900..1a1a0d825 100644
--- a/tools/build/defs.bzl
+++ b/tools/build/defs.bzl
@@ -8,6 +8,7 @@ load("@rules_pkg//:pkg.bzl", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar")
 load("@io_bazel_rules_docker//go:image.bzl", _go_image = "go_image")
 load("@io_bazel_rules_docker//container:container.bzl", _container_image = "container_image")
 load("@pydeps//:requirements.bzl", _py_requirement = "requirement")
+load("//tools/build:tags.bzl", _go_suffixes = "go_suffixes")
 
 container_image = _container_image
 cc_binary = _cc_binary
@@ -18,6 +19,7 @@ cc_test = _cc_test
 cc_toolchain = "@bazel_tools//tools/cpp:current_cc_toolchain"
 go_image = _go_image
 go_embed_data = _go_embed_data
+go_suffixes = _go_suffixes
 gtest = "@com_google_googletest//:gtest"
 loopback = "//tools/build:loopback"
 proto_library = native.proto_library
diff --git a/tools/build/tags.bzl b/tools/build/tags.bzl
new file mode 100644
index 000000000..e99c87f81
--- /dev/null
+++ b/tools/build/tags.bzl
@@ -0,0 +1,36 @@
+"""List of special Go suffixes."""
+
+go_suffixes = [
+    "_386",
+    "_386_unsafe",
+    "_amd64",
+    "_amd64_unsafe",
+    "_aarch64",
+    "_aarch64_unsafe",
+    "_arm",
+    "_arm_unsafe",
+    "_arm64",
+    "_arm64_unsafe",
+    "_mips",
+    "_mips_unsafe",
+    "_mipsle",
+    "_mipsle_unsafe",
+    "_mips64",
+    "_mips64_unsafe",
+    "_mips64le",
+    "_mips64le_unsafe",
+    "_ppc64",
+    "_ppc64_unsafe",
+    "_ppc64le",
+    "_ppc64le_unsafe",
+    "_riscv64",
+    "_riscv64_unsafe",
+    "_s390x",
+    "_s390x_unsafe",
+    "_sparc64",
+    "_sparc64_unsafe",
+    "_wasm",
+    "_wasm_unsafe",
+    "_linux",
+    "_linux_unsafe",
+]
diff --git a/tools/defs.bzl b/tools/defs.bzl
index ce677cbbf..5d5fa134a 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/build:defs.bzl", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _container_image = "container_image", _default_installer = "default_installer", _default_net_util = "default_net_util", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_image = "go_image", _go_library = "go_library", _go_proto_library = "go_proto_library", _go_test = "go_test", _go_tool_library = "go_tool_library", _gtest = "gtest", _loopback = "loopback", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar", _proto_library = "proto_library", _py_binary = "py_binary", _py_library = "py_library", _py_requirement = "py_requirement", _py_test = "py_test", _select_arch = "select_arch", _select_system = "select_system")
+load("//tools/build:defs.bzl", "go_suffixes", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _container_image = "container_image", _default_installer = "default_installer", _default_net_util = "default_net_util", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_image = "go_image", _go_library = "go_library", _go_proto_library = "go_proto_library", _go_test = "go_test", _go_tool_library = "go_tool_library", _gtest = "gtest", _loopback = "loopback", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar", _proto_library = "proto_library", _py_binary = "py_binary", _py_library = "py_library", _py_requirement = "py_requirement", _py_test = "py_test", _select_arch = "select_arch", _select_system = "select_system")
 
 # Delegate directly.
 cc_binary = _cc_binary
@@ -45,6 +45,34 @@ def go_binary(name, **kwargs):
         **kwargs
     )
 
+def calculate_sets(srcs):
+    """Calculates special Go sets for templates.
+
+    Args:
+      srcs: the full set of Go sources.
+
+    Returns:
+      A dictionary of the form:
+
+      "": [src1.go, src2.go]
+      "suffix": [src3suffix.go, src4suffix.go]
+
+      Note that suffix will typically start with '_'.
+    """
+    result = dict()
+    for file in srcs:
+        if not file.endswith(".go"):
+            continue
+        target = ""
+        for suffix in go_suffixes:
+            if file.endswith(suffix + ".go"):
+                target = suffix
+        if not target in result:
+            result[target] = [file]
+        else:
+            result[target].append(file)
+    return result
+
 def go_library(name, srcs, deps = [], imports = [], stateify = True, marshal = False, **kwargs):
     """Wraps the standard go_library and does stateification and marshalling.
 
@@ -70,39 +98,49 @@ def go_library(name, srcs, deps = [], imports = [], stateify = True, marshal = F
       marshal: whether marshal is enabled (default: false).
       **kwargs: standard go_library arguments.
     """
+    all_srcs = srcs
+    all_deps = deps
     if stateify:
         # Only do stateification for non-state packages without manual autogen.
-        go_stateify(
-            name = name + "_state_autogen",
-            srcs = [src for src in srcs if src.endswith(".go")],
-            imports = imports,
-            package = name,
-            arch = select_arch(),
-            out = name + "_state_autogen.go",
-        )
-        all_srcs = srcs + [name + "_state_autogen.go"]
-        if "//pkg/state" not in deps:
-            all_deps = deps + ["//pkg/state"]
-        else:
-            all_deps = deps
-    else:
-        all_deps = deps
-        all_srcs = srcs
+        # First, we need to segregate the input files via the special suffixes,
+        # and calculate the final output set.
+        state_sets = calculate_sets(srcs)
+        for (suffix, srcs) in state_sets.items():
+            go_stateify(
+                name = name + suffix + "_state_autogen",
+                srcs = srcs,
+                imports = imports,
+                package = name,
+                out = name + suffix + "_state_autogen.go",
+            )
+        all_srcs = all_srcs + [
+            name + suffix + "_state_autogen.go"
+            for suffix in state_sets.keys()
+        ]
+        if "//pkg/state" not in all_deps:
+            all_deps = all_deps + ["//pkg/state"]
+
     if marshal:
-        go_marshal(
-            name = name + "_abi_autogen",
-            srcs = [src for src in srcs if src.endswith(".go")],
-            debug = False,
-            imports = imports,
-            package = name,
-        )
+        # See above.
+        marshal_sets = calculate_sets(srcs)
+        for (suffix, srcs) in marshal_sets.items():
+            go_marshal(
+                name = name + suffix + "_abi_autogen",
+                srcs = srcs,
+                debug = False,
+                imports = imports,
+                package = name,
+            )
         extra_deps = [
             dep
             for dep in marshal_deps
             if not dep in all_deps
         ]
         all_deps = all_deps + extra_deps
-        all_srcs = srcs + [name + "_abi_autogen_unsafe.go"]
+        all_srcs = all_srcs + [
+            name + suffix + "_abi_autogen_unsafe.go"
+            for suffix in marshal_sets.keys()
+        ]
 
     _go_library(
         name = name,
@@ -115,13 +153,16 @@ def go_library(name, srcs, deps = [], imports = [], stateify = True, marshal = F
         # Ignore importpath for go_test.
         kwargs.pop("importpath", None)
 
-        _go_test(
-            name = name + "_abi_autogen_test",
-            srcs = [name + "_abi_autogen_test.go"],
-            library = ":" + name,
-            deps = marshal_test_deps,
-            **kwargs
-        )
+        # See above.
+        marshal_sets = calculate_sets(srcs)
+        for (suffix, srcs) in marshal_sets.items():
+            _go_test(
+                name = name + suffix + "_abi_autogen_test",
+                srcs = [name + suffix + "_abi_autogen_test.go"],
+                library = ":" + name + suffix,
+                deps = marshal_test_deps,
+                **kwargs
+            )
 
 def proto_library(name, srcs, **kwargs):
     """Wraps the standard proto_library.
diff --git a/tools/go_marshal/gomarshal/BUILD b/tools/go_marshal/gomarshal/BUILD
index c92b59dd6..b5d5a4487 100644
--- a/tools/go_marshal/gomarshal/BUILD
+++ b/tools/go_marshal/gomarshal/BUILD
@@ -14,4 +14,5 @@ go_library(
     visibility = [
         "//:sandbox",
     ],
+    deps = ["//tools/tags"],
 )
diff --git a/tools/go_marshal/gomarshal/generator.go b/tools/go_marshal/gomarshal/generator.go
index af90bdecb..0b3f600fe 100644
--- a/tools/go_marshal/gomarshal/generator.go
+++ b/tools/go_marshal/gomarshal/generator.go
@@ -23,6 +23,9 @@ import (
 	"go/token"
 	"os"
 	"sort"
+	"strings"
+
+	"gvisor.dev/gvisor/tools/tags"
 )
 
 const (
@@ -104,6 +107,14 @@ func NewGenerator(srcs []string, out, outTest, pkg string, imports []string) (*G
 func (g *Generator) writeHeader() error {
 	var b sourceBuffer
 	b.emit("// Automatically generated marshal implementation. See tools/go_marshal.\n\n")
+
+	// Emit build tags.
+	if t := tags.Aggregate(g.inputs); len(t) > 0 {
+		b.emit(strings.Join(t.Lines(), "\n"))
+		b.emit("\n")
+	}
+
+	// Package header.
 	b.emit("package %s\n\n", g.pkg)
 	if err := b.write(g.output); err != nil {
 		return err
diff --git a/tools/go_stateify/BUILD b/tools/go_stateify/BUILD
index a133d6f8b..6036faf7b 100644
--- a/tools/go_stateify/BUILD
+++ b/tools/go_stateify/BUILD
@@ -6,4 +6,5 @@ go_binary(
     name = "stateify",
     srcs = ["main.go"],
     visibility = ["//visibility:public"],
+    deps = ["//tools/tags"],
 )
diff --git a/tools/go_stateify/defs.bzl b/tools/go_stateify/defs.bzl
index 0f261d89f..bdb966362 100644
--- a/tools/go_stateify/defs.bzl
+++ b/tools/go_stateify/defs.bzl
@@ -7,7 +7,6 @@ def _go_stateify_impl(ctx):
     # Run the stateify command.
     args = ["-output=%s" % output.path]
     args.append("-pkg=%s" % ctx.attr.package)
-    args.append("-arch=%s" % ctx.attr.arch)
     if ctx.attr._statepkg:
         args.append("-statepkg=%s" % ctx.attr._statepkg)
     if ctx.attr.imports:
@@ -47,15 +46,8 @@ for statified types.
             doc = "The package name for the input sources.",
             mandatory = True,
         ),
-        "arch": attr.string(
-            doc = "Target platform.",
-            mandatory = True,
-        ),
         "out": attr.output(
-            doc = """
-The name of the generated file output. This must not conflict with any other
-files and must be added to the srcs of the relevant go_library.
-""",
+            doc = "Name of the generator output file.",
             mandatory = True,
         ),
         "_tool": attr.label(
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index 7d5d291e6..aa9d4543e 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -22,12 +22,12 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
-	"io/ioutil"
 	"os"
-	"path/filepath"
 	"reflect"
 	"strings"
 	"sync"
+
+	"gvisor.dev/gvisor/tools/tags"
 )
 
 var (
@@ -35,113 +35,8 @@ var (
 	imports  = flag.String("imports", "", "extra imports for the output file")
 	output   = flag.String("output", "", "output file")
 	statePkg = flag.String("statepkg", "", "state import package; defaults to empty")
-	arch     = flag.String("arch", "", "specify the target platform")
 )
 
-// The known architectures.
-var okgoarch = []string{
-	"386",
-	"amd64",
-	"arm",
-	"arm64",
-	"mips",
-	"mipsle",
-	"mips64",
-	"mips64le",
-	"ppc64",
-	"ppc64le",
-	"riscv64",
-	"s390x",
-	"sparc64",
-	"wasm",
-}
-
-// readfile returns the content of the named file.
-func readfile(file string) string {
-	data, err := ioutil.ReadFile(file)
-	if err != nil {
-		panic(fmt.Sprintf("readfile err: %v", err))
-	}
-	return string(data)
-}
-
-// matchfield reports whether the field (x,y,z) matches this build.
-// all the elements in the field must be satisfied.
-func matchfield(f string, goarch string) bool {
-	for _, tag := range strings.Split(f, ",") {
-		if !matchtag(tag, goarch) {
-			return false
-		}
-	}
-	return true
-}
-
-// matchtag reports whether the tag (x or !x) matches this build.
-func matchtag(tag string, goarch string) bool {
-	if tag == "" {
-		return false
-	}
-	if tag[0] == '!' {
-		if len(tag) == 1 || tag[1] == '!' {
-			return false
-		}
-		return !matchtag(tag[1:], goarch)
-	}
-	return tag == goarch
-}
-
-// canBuild reports whether we can build this file for target platform by
-// checking file name and build tags. The code is derived from the Go source
-// cmd.dist.build.shouldbuild.
-func canBuild(file, goTargetArch string) bool {
-	name := filepath.Base(file)
-	excluded := func(list []string, ok string) bool {
-		for _, x := range list {
-			if x == ok || (ok == "android" && x == "linux") || (ok == "illumos" && x == "solaris") {
-				continue
-			}
-			i := strings.Index(name, x)
-			if i <= 0 || name[i-1] != '_' {
-				continue
-			}
-			i += len(x)
-			if i == len(name) || name[i] == '.' || name[i] == '_' {
-				return true
-			}
-		}
-		return false
-	}
-	if excluded(okgoarch, goTargetArch) {
-		return false
-	}
-
-	// Check file contents for // +build lines.
-	for _, p := range strings.Split(readfile(file), "\n") {
-		p = strings.TrimSpace(p)
-		if p == "" {
-			continue
-		}
-		if !strings.HasPrefix(p, "//") {
-			break
-		}
-		if !strings.Contains(p, "+build") {
-			continue
-		}
-		fields := strings.Fields(p[2:])
-		if len(fields) < 1 || fields[0] != "+build" {
-			continue
-		}
-		for _, p := range fields[1:] {
-			if matchfield(p, goTargetArch) {
-				goto fieldmatch
-			}
-		}
-		return false
-	fieldmatch:
-	}
-	return true
-}
-
 // resolveTypeName returns a qualified type name.
 func resolveTypeName(name string, typ ast.Expr) (field string, qualified string) {
 	for done := false; !done; {
@@ -329,8 +224,15 @@ func main() {
 		fmt.Fprintf(outputFile, "	m.Save(\"%s\", &x.%s)\n", name, name)
 	}
 
-	// Emit the package name.
+	// Automated warning.
 	fmt.Fprint(outputFile, "// automatically generated by stateify.\n\n")
+
+	// Emit build tags.
+	if t := tags.Aggregate(flag.Args()); len(t) > 0 {
+		fmt.Fprintf(outputFile, "%s\n\n", strings.Join(t.Lines(), "\n"))
+	}
+
+	// Emit the package name.
 	fmt.Fprintf(outputFile, "package %s\n\n", *pkg)
 
 	// Emit the imports lazily.
@@ -364,10 +266,6 @@ func main() {
 			os.Exit(1)
 		}
 
-		if !canBuild(filename, *arch) {
-			continue
-		}
-
 		files = append(files, f)
 	}
 
diff --git a/tools/tags/BUILD b/tools/tags/BUILD
new file mode 100644
index 000000000..1c02e2c89
--- /dev/null
+++ b/tools/tags/BUILD
@@ -0,0 +1,11 @@
+load("//tools:defs.bzl", "go_library")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "tags",
+    srcs = ["tags.go"],
+    marshal = False,
+    stateify = False,
+    visibility = ["//tools:__subpackages__"],
+)
diff --git a/tools/tags/tags.go b/tools/tags/tags.go
new file mode 100644
index 000000000..f35904e0a
--- /dev/null
+++ b/tools/tags/tags.go
@@ -0,0 +1,89 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package tags is a utility for parsing build tags.
+package tags
+
+import (
+	"fmt"
+	"io/ioutil"
+	"strings"
+)
+
+// OrSet is a set of tags on a single line.
+//
+// Note that tags may include ",", and we don't distinguish this case in the
+// logic below. Ideally, this constraints can be split into separate top-level
+// build tags in order to resolve any issues.
+type OrSet []string
+
+// Line returns the line for this or.
+func (or OrSet) Line() string {
+	return fmt.Sprintf("// +build %s", strings.Join([]string(or), " "))
+}
+
+// AndSet is the set of all OrSets.
+type AndSet []OrSet
+
+// Lines returns the lines to be printed.
+func (and AndSet) Lines() (ls []string) {
+	for _, or := range and {
+		ls = append(ls, or.Line())
+	}
+	return
+}
+
+// Join joins this AndSet with another.
+func (and AndSet) Join(other AndSet) AndSet {
+	return append(and, other...)
+}
+
+// Tags returns the unique set of +build tags.
+//
+// Derived form the runtime's canBuild.
+func Tags(file string) (tags AndSet) {
+	data, err := ioutil.ReadFile(file)
+	if err != nil {
+		return nil
+	}
+	// Check file contents for // +build lines.
+	for _, p := range strings.Split(string(data), "\n") {
+		p = strings.TrimSpace(p)
+		if p == "" {
+			continue
+		}
+		if !strings.HasPrefix(p, "//") {
+			break
+		}
+		if !strings.Contains(p, "+build") {
+			continue
+		}
+		fields := strings.Fields(p[2:])
+		if len(fields) < 1 || fields[0] != "+build" {
+			continue
+		}
+		tags = append(tags, OrSet(fields[1:]))
+	}
+	return tags
+}
+
+// Aggregate aggregates all tags from a set of files.
+//
+// Note that these may be in conflict, in which case the build will fail.
+func Aggregate(files []string) (tags AndSet) {
+	for _, file := range files {
+		tags = tags.Join(Tags(file))
+	}
+	return tags
+}