1 files changed, 18 insertions, 6 deletions
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index db7a7107b..309ee9c21 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -23,13 +23,16 @@ import (
 	"go/parser"
 	"go/token"
 	"os"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"sync"
+
+	"gvisor.dev/gvisor/tools/tags"
 )
 
 var (
-	pkg      = flag.String("pkg", "", "output package")
+	fullPkg  = flag.String("fullpkg", "", "fully qualified output package")
 	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")
@@ -168,7 +171,7 @@ func main() {
 		flag.Usage()
 		os.Exit(1)
 	}
-	if *pkg == "" {
+	if *fullPkg == "" {
 		fmt.Fprintf(os.Stderr, "Error: package required.")
 		os.Exit(1)
 	}
@@ -200,10 +203,10 @@ func main() {
 
 	// Declare our emission closures.
 	emitRegister := func(name string) {
-		initCalls = append(initCalls, fmt.Sprintf("%sRegister(\"%s.%s\", (*%s)(nil), state.Fns{Save: (*%s).save, Load: (*%s).load})", statePrefix, *pkg, name, name, name, name))
+		initCalls = append(initCalls, fmt.Sprintf("%sRegister(\"%s.%s\", (*%s)(nil), state.Fns{Save: (*%s).save, Load: (*%s).load})", statePrefix, *fullPkg, name, name, name, name))
 	}
 	emitZeroCheck := func(name string) {
-		fmt.Fprintf(outputFile, "	if !%sIsZeroValue(x.%s) { m.Failf(\"%s is %%v, expected zero\", x.%s) }\n", statePrefix, name, name, name)
+		fmt.Fprintf(outputFile, "	if !%sIsZeroValue(&x.%s) { m.Failf(\"%s is %%#v, expected zero\", &x.%s) }\n", statePrefix, name, name, name)
 	}
 	emitLoadValue := func(name, typName string) {
 		fmt.Fprintf(outputFile, "	m.LoadValue(\"%s\", new(%s), func(y interface{}) { x.load%s(y.(%s)) })\n", name, typName, camelCased(name), typName)
@@ -222,9 +225,17 @@ 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")
-	fmt.Fprintf(outputFile, "package %s\n\n", *pkg)
+
+	// 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.
+	_, pkg := filepath.Split(*fullPkg)
+	fmt.Fprintf(outputFile, "package %s\n\n", pkg)
 
 	// Emit the imports lazily.
 	var once sync.Once
@@ -256,6 +267,7 @@ func main() {
 			fmt.Fprintf(os.Stderr, "Input %q can't be parsed: %v\n", filename, err)
 			os.Exit(1)
 		}
+
 		files = append(files, f)
 	}