stateify: support explicit annotation mode; convert refs and stack packages.

We have been unnecessarily creating too many savable types implicitly.

PiperOrigin-RevId: 206334201
Change-Id: Idc5a3a14bfb7ee125c4f2bb2b1c53164e46f29a8

1 files changed, 51 insertions, 15 deletions

diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index 6c3583c62..231c6d80b 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -25,6 +25,7 @@ import (
 	"os"
 	"reflect"
 	"strings"
+	"sync"
 )
 
 var (
@@ -32,6 +33,7 @@ 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")
+	explicit = flag.Bool("explicit", false, "only generate for types explicitly tagged '// +stateify savable'")
 )
 
 // resolveTypeName returns a qualified type name.
@@ -224,16 +226,24 @@ func main() {
 	// Emit the package name.
 	fmt.Fprint(outputFile, "// automatically generated by stateify.\n\n")
 	fmt.Fprintf(outputFile, "package %s\n\n", *pkg)
-	fmt.Fprint(outputFile, "import (\n")
-	if *statePkg != "" {
-		fmt.Fprintf(outputFile, "	\"%s\"\n", *statePkg)
-	}
-	if *imports != "" {
-		for _, i := range strings.Split(*imports, ",") {
-			fmt.Fprintf(outputFile, "	\"%s\"\n", i)
-		}
+
+	// Emit the imports lazily.
+	var once sync.Once
+	maybeEmitImports := func() {
+		once.Do(func() {
+			// Emit the imports.
+			fmt.Fprint(outputFile, "import (\n")
+			if *statePkg != "" {
+				fmt.Fprintf(outputFile, "	\"%s\"\n", *statePkg)
+			}
+			if *imports != "" {
+				for _, i := range strings.Split(*imports, ",") {
+					fmt.Fprintf(outputFile, "	\"%s\"\n", i)
+				}
+			}
+			fmt.Fprint(outputFile, ")\n\n")
+		})
 	}
-	fmt.Fprint(outputFile, ")\n\n")
 
 	files := make([]*ast.File, 0, len(flag.Args()))
 
@@ -241,7 +251,7 @@ func main() {
 	for _, filename := range flag.Args() {
 		// Parse the file.
 		fset := token.NewFileSet()
-		f, err := parser.ParseFile(fset, filename, nil, 0)
+		f, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
 		if err != nil {
 			// Not a valid input file?
 			fmt.Fprintf(os.Stderr, "Input %q can't be parsed: %v\n", filename, err)
@@ -308,6 +318,26 @@ func main() {
 				continue
 			}
 
+			if *explicit {
+				// In explicit mode, only generate code for
+				// types explicitly marked
+				// "// +stateify savable" in one of the
+				// proceeding comment lines.
+				if d.Doc == nil {
+					continue
+				}
+				savable := false
+				for _, l := range d.Doc.List {
+					if l.Text == "// +stateify savable" {
+						savable = true
+						break
+					}
+				}
+				if !savable {
+					continue
+				}
+			}
+
 			for _, gs := range d.Specs {
 				ts := gs.(*ast.TypeSpec)
 				switch ts.Type.(type) {
@@ -315,6 +345,8 @@ func main() {
 					// Don't register.
 					break
 				case *ast.StructType:
+					maybeEmitImports()
+
 					ss := ts.Type.(*ast.StructType)
 
 					// Define beforeSave if a definition was not found. This
@@ -360,6 +392,8 @@ func main() {
 					// Add to our registration.
 					emitRegister(ts.Name.Name)
 				case *ast.Ident, *ast.SelectorExpr, *ast.ArrayType:
+					maybeEmitImports()
+
 					_, val := resolveTypeName(ts.Name.Name, ts.Type)
 
 					// Dispatch directly.
@@ -377,10 +411,12 @@ func main() {
 		}
 	}
 
-	// Emit the init() function.
-	fmt.Fprintf(outputFile, "func init() {\n")
-	for _, ic := range initCalls {
-		fmt.Fprintf(outputFile, "	%s\n", ic)
+	if len(initCalls) > 0 {
+		// Emit the init() function.
+		fmt.Fprintf(outputFile, "func init() {\n")
+		for _, ic := range initCalls {
+			fmt.Fprintf(outputFile, "	%s\n", ic)
+		}
+		fmt.Fprintf(outputFile, "}\n")
 	}
-	fmt.Fprintf(outputFile, "}\n")
 }