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

3 files changed, 12 insertions, 121 deletions

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)
 	}