summaryrefslogtreecommitdiffhomepage
path: root/tools/go_stateify
diff options
context:
space:
mode:
Diffstat (limited to 'tools/go_stateify')
-rw-r--r--tools/go_stateify/BUILD9
-rw-r--r--tools/go_stateify/defs.bzl77
-rw-r--r--tools/go_stateify/main.go386
3 files changed, 472 insertions, 0 deletions
diff --git a/tools/go_stateify/BUILD b/tools/go_stateify/BUILD
new file mode 100644
index 000000000..edbeb4e2d
--- /dev/null
+++ b/tools/go_stateify/BUILD
@@ -0,0 +1,9 @@
+package(licenses = ["notice"]) # Apache 2.0
+
+load("@io_bazel_rules_go//go:def.bzl", "go_binary")
+
+go_binary(
+ name = "stateify",
+ srcs = ["main.go"],
+ visibility = ["//visibility:public"],
+)
diff --git a/tools/go_stateify/defs.bzl b/tools/go_stateify/defs.bzl
new file mode 100644
index 000000000..87fdc0d28
--- /dev/null
+++ b/tools/go_stateify/defs.bzl
@@ -0,0 +1,77 @@
+"""Stateify is a tool for generating state wrappers for Go types.
+
+The go_stateify rule is used to generate a file that will appear in a Go
+target; the output file should appear explicitly in a srcs list. For example:
+
+go_stateify(
+ name = "foo_state",
+ srcs = ["foo.go"],
+ out = "foo_state.go",
+ package = "foo",
+)
+
+go_library(
+ name = "foo",
+ srcs = [
+ "foo.go",
+ "foo_state.go",
+ ],
+ deps = [
+ "//pkg/state",
+ ],
+)
+"""
+
+def _go_stateify_impl(ctx):
+ """Implementation for the stateify tool."""
+ output = ctx.outputs.out
+
+ # Run the stateify command.
+ args = ["-output=%s" % output.path]
+ args += ["-pkg=%s" % ctx.attr.package]
+ if ctx.attr._statepkg:
+ args += ["-statepkg=%s" % ctx.attr._statepkg]
+ if ctx.attr.imports:
+ args += ["-imports=%s" % ",".join(ctx.attr.imports)]
+ args += ["--"]
+ for src in ctx.attr.srcs:
+ args += [f.path for f in src.files]
+ ctx.actions.run(
+ inputs = ctx.files.srcs,
+ outputs = [output],
+ mnemonic = "GoStateify",
+ progress_message = "Generating state library %s" % ctx.label,
+ arguments = args,
+ executable = ctx.executable._tool,
+ )
+
+"""
+Generates save and restore logic from a set of Go files.
+
+
+Args:
+ name: the name of the rule.
+ srcs: the input source files. These files should include all structs in the package that need to be saved.
+ imports: an optional list of extra non-aliased, Go-style absolute import paths.
+ out: 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.
+ package: the package name for the input sources.
+"""
+
+go_stateify = rule(
+ attrs = {
+ "srcs": attr.label_list(
+ mandatory = True,
+ allow_files = True,
+ ),
+ "imports": attr.string_list(mandatory = False),
+ "package": attr.string(mandatory = True),
+ "out": attr.output(mandatory = True),
+ "_tool": attr.label(
+ executable = True,
+ cfg = "host",
+ default = Label("//tools/go_stateify:stateify"),
+ ),
+ "_statepkg": attr.string(default = "gvisor.googlesource.com/gvisor/pkg/state"),
+ },
+ implementation = _go_stateify_impl,
+)
diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
new file mode 100644
index 000000000..5eb4fe51f
--- /dev/null
+++ b/tools/go_stateify/main.go
@@ -0,0 +1,386 @@
+// Copyright 2018 Google Inc.
+//
+// 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.
+
+// Stateify provides a simple way to generate Load/Save methods based on
+// existing types and struct tags.
+package main
+
+import (
+ "flag"
+ "fmt"
+ "go/ast"
+ "go/parser"
+ "go/token"
+ "os"
+ "reflect"
+ "strings"
+)
+
+var (
+ pkg = flag.String("pkg", "", "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")
+)
+
+// resolveTypeName returns a qualified type name.
+func resolveTypeName(name string, typ ast.Expr) (field string, qualified string) {
+ for done := false; !done; {
+ // Resolve star expressions.
+ switch rs := typ.(type) {
+ case *ast.StarExpr:
+ qualified += "*"
+ typ = rs.X
+ case *ast.ArrayType:
+ if rs.Len == nil {
+ // Slice type declaration.
+ qualified += "[]"
+ } else {
+ // Array type declaration.
+ qualified += "[" + rs.Len.(*ast.BasicLit).Value + "]"
+ }
+ typ = rs.Elt
+ default:
+ // No more descent.
+ done = true
+ }
+ }
+
+ // Resolve a package selector.
+ sel, ok := typ.(*ast.SelectorExpr)
+ if ok {
+ qualified = qualified + sel.X.(*ast.Ident).Name + "."
+ typ = sel.Sel
+ }
+
+ // Figure out actual type name.
+ ident, ok := typ.(*ast.Ident)
+ if !ok {
+ panic(fmt.Sprintf("type not supported: %s (involves anonymous types?)", name))
+ }
+ field = ident.Name
+ qualified = qualified + field
+ return
+}
+
+// extractStateTag pulls the relevant state tag.
+func extractStateTag(tag *ast.BasicLit) string {
+ if tag == nil {
+ return ""
+ }
+ if len(tag.Value) < 2 {
+ return ""
+ }
+ return reflect.StructTag(tag.Value[1 : len(tag.Value)-1]).Get("state")
+}
+
+// scanFunctions is a set of functions passed to scanFields.
+type scanFunctions struct {
+ zerovalue func(name string)
+ normal func(name string)
+ wait func(name string)
+ value func(name, typName string)
+}
+
+// scanFields scans the fields of a struct.
+//
+// Each provided function will be applied to appropriately tagged fields, or
+// skipped if nil.
+//
+// Fields tagged nosave are skipped.
+func scanFields(ss *ast.StructType, fn scanFunctions) {
+ if ss.Fields.List == nil {
+ // No fields.
+ return
+ }
+
+ // Scan all fields.
+ for _, field := range ss.Fields.List {
+ // Calculate the name.
+ name := ""
+ if field.Names != nil {
+ // It's a named field; override.
+ name = field.Names[0].Name
+ } else {
+ // Anonymous types can't be embedded, so we don't need
+ // to worry about providing a useful name here.
+ name, _ = resolveTypeName("", field.Type)
+ }
+
+ // Skip _ fields.
+ if name == "_" {
+ continue
+ }
+
+ switch tag := extractStateTag(field.Tag); tag {
+ case "zerovalue":
+ if fn.zerovalue != nil {
+ fn.zerovalue(name)
+ }
+
+ case "":
+ if fn.normal != nil {
+ fn.normal(name)
+ }
+
+ case "wait":
+ if fn.wait != nil {
+ fn.wait(name)
+ }
+
+ case "manual", "nosave", "ignore":
+ // Do nothing.
+
+ default:
+ if strings.HasPrefix(tag, ".(") && strings.HasSuffix(tag, ")") {
+ if fn.value != nil {
+ fn.value(name, tag[2:len(tag)-1])
+ }
+ }
+ }
+ }
+}
+
+func camelCased(name string) string {
+ return strings.ToUpper(name[:1]) + name[1:]
+}
+
+func main() {
+ // Parse flags.
+ flag.Usage = func() {
+ fmt.Fprintf(os.Stderr, "Usage: %s [options]\n", os.Args[0])
+ flag.PrintDefaults()
+ }
+ flag.Parse()
+ if len(flag.Args()) == 0 {
+ flag.Usage()
+ os.Exit(1)
+ }
+ if *pkg == "" {
+ fmt.Fprintf(os.Stderr, "Error: package required.")
+ os.Exit(1)
+ }
+
+ // Open the output file.
+ var (
+ outputFile *os.File
+ err error
+ )
+ if *output == "" || *output == "-" {
+ outputFile = os.Stdout
+ } else {
+ outputFile, err = os.OpenFile(*output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, "Error opening output %q: %v", *output, err)
+ }
+ defer outputFile.Close()
+ }
+
+ // Set the statePrefix for below, depending on the import.
+ statePrefix := ""
+ if *statePkg != "" {
+ parts := strings.Split(*statePkg, "/")
+ statePrefix = parts[len(parts)-1] + "."
+ }
+
+ // initCalls is dumped at the end.
+ var initCalls []string
+
+ // 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))
+ }
+ 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)
+ }
+ 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)
+ }
+ emitLoad := func(name string) {
+ fmt.Fprintf(outputFile, " m.Load(\"%s\", &x.%s)\n", name, name)
+ }
+ emitLoadWait := func(name string) {
+ fmt.Fprintf(outputFile, " m.LoadWait(\"%s\", &x.%s)\n", name, name)
+ }
+ emitSaveValue := func(name, typName string) {
+ fmt.Fprintf(outputFile, " var %s %s = x.save%s()\n", name, typName, camelCased(name))
+ fmt.Fprintf(outputFile, " m.SaveValue(\"%s\", %s)\n", name, name)
+ }
+ emitSave := func(name string) {
+ fmt.Fprintf(outputFile, " m.Save(\"%s\", &x.%s)\n", name, name)
+ }
+
+ // Emit the package name.
+ fmt.Fprintf(outputFile, "// automatically generated by stateify.\n\n")
+ fmt.Fprintf(outputFile, "package %s\n\n", *pkg)
+ fmt.Fprintln(outputFile, "import (")
+ if *statePkg != "" {
+ fmt.Fprintf(outputFile, " \"%s\"\n", *statePkg)
+ }
+ if *imports != "" {
+ for _, i := range strings.Split(*imports, ",") {
+ fmt.Fprintf(outputFile, " \"%s\"\n", i)
+ }
+ }
+ fmt.Fprintln(outputFile, ")\n")
+
+ files := make([]*ast.File, 0, len(flag.Args()))
+
+ // Parse the input files.
+ for _, filename := range flag.Args() {
+ // Parse the file.
+ fset := token.NewFileSet()
+ f, err := parser.ParseFile(fset, filename, nil, 0)
+ if err != nil {
+ // Not a valid input file?
+ fmt.Fprintf(os.Stderr, "Input %q can't be parsed: %v\n", filename, err)
+ os.Exit(1)
+ }
+ files = append(files, f)
+ }
+
+ type method struct {
+ receiver string
+ name string
+ }
+
+ // Search for and add all methods with a pointer receiver and no other
+ // arguments to a set. We support auto-detecting the existence of
+ // several different methods with this signature.
+ simpleMethods := map[method]struct{}{}
+ for _, f := range files {
+
+ // Go over all functions.
+ for _, decl := range f.Decls {
+ d, ok := decl.(*ast.FuncDecl)
+ if !ok {
+ continue
+ }
+ if d.Name == nil || d.Recv == nil || d.Type == nil {
+ // Not a named method.
+ continue
+ }
+ if len(d.Recv.List) != 1 {
+ // Wrong number of receivers?
+ continue
+ }
+ if d.Type.Params != nil && len(d.Type.Params.List) != 0 {
+ // Has argument(s).
+ continue
+ }
+ if d.Type.Results != nil && len(d.Type.Results.List) != 0 {
+ // Has return(s).
+ continue
+ }
+
+ pt, ok := d.Recv.List[0].Type.(*ast.StarExpr)
+ if !ok {
+ // Not a pointer receiver.
+ continue
+ }
+
+ t, ok := pt.X.(*ast.Ident)
+ if !ok {
+ // This shouldn't happen with valid Go.
+ continue
+ }
+
+ simpleMethods[method{t.Name, d.Name.Name}] = struct{}{}
+ }
+ }
+
+ for _, f := range files {
+ // Go over all named types.
+ for _, decl := range f.Decls {
+ d, ok := decl.(*ast.GenDecl)
+ if !ok || d.Tok != token.TYPE {
+ continue
+ }
+
+ for _, gs := range d.Specs {
+ ts := gs.(*ast.TypeSpec)
+ switch ts.Type.(type) {
+ case *ast.InterfaceType, *ast.ChanType, *ast.FuncType, *ast.ParenExpr, *ast.StarExpr:
+ // Don't register.
+ break
+ case *ast.StructType:
+ ss := ts.Type.(*ast.StructType)
+
+ // Define beforeSave if a definition was not found. This
+ // prevents the code from compiling if a custom beforeSave
+ // was defined in a file not provided to this binary and
+ // prevents inherited methods from being called multiple times
+ // by overriding them.
+ if _, ok := simpleMethods[method{ts.Name.Name, "beforeSave"}]; !ok {
+ fmt.Fprintf(outputFile, "func (x *%s) beforeSave() {}\n", ts.Name.Name)
+ }
+
+ // Generate the save method.
+ fmt.Fprintf(outputFile, "func (x *%s) save(m %sMap) {\n", ts.Name.Name, statePrefix)
+ fmt.Fprintf(outputFile, " x.beforeSave()\n")
+ scanFields(ss, scanFunctions{zerovalue: emitZeroCheck})
+ scanFields(ss, scanFunctions{value: emitSaveValue})
+ scanFields(ss, scanFunctions{normal: emitSave, wait: emitSave})
+ fmt.Fprintf(outputFile, "}\n\n")
+
+ // Define afterLoad if a definition was not found. We do this
+ // for the same reason that we do it for beforeSave.
+ _, hasAfterLoad := simpleMethods[method{ts.Name.Name, "afterLoad"}]
+ if !hasAfterLoad {
+ fmt.Fprintf(outputFile, "func (x *%s) afterLoad() {}\n", ts.Name.Name)
+ }
+
+ // Generate the load method.
+ //
+ // Note that the manual loads always follow the
+ // automated loads.
+ fmt.Fprintf(outputFile, "func (x *%s) load(m %sMap) {\n", ts.Name.Name, statePrefix)
+ scanFields(ss, scanFunctions{normal: emitLoad, wait: emitLoadWait})
+ scanFields(ss, scanFunctions{value: emitLoadValue})
+ if hasAfterLoad {
+ // The call to afterLoad is made conditionally, because when
+ // AfterLoad is called, the object encodes a dependency on
+ // referred objects (i.e. fields). This means that afterLoad
+ // will not be called until the other afterLoads are called.
+ fmt.Fprintf(outputFile, " m.AfterLoad(x.afterLoad)\n")
+ }
+ fmt.Fprintf(outputFile, "}\n\n")
+
+ // Add to our registration.
+ emitRegister(ts.Name.Name)
+ case *ast.Ident, *ast.SelectorExpr, *ast.ArrayType:
+ _, val := resolveTypeName(ts.Name.Name, ts.Type)
+
+ // Dispatch directly.
+ fmt.Fprintf(outputFile, "func (x *%s) save(m %sMap) {\n", ts.Name.Name, statePrefix)
+ fmt.Fprintf(outputFile, " m.SaveValue(\"\", (%s)(*x))\n", val)
+ fmt.Fprintf(outputFile, "}\n\n")
+ fmt.Fprintf(outputFile, "func (x *%s) load(m %sMap) {\n", ts.Name.Name, statePrefix)
+ fmt.Fprintf(outputFile, " m.LoadValue(\"\", new(%s), func(y interface{}) { *x = (%s)(y.(%s)) })\n", val, ts.Name.Name, val)
+ fmt.Fprintf(outputFile, "}\n\n")
+
+ // See above.
+ emitRegister(ts.Name.Name)
+ }
+ }
+ }
+ }
+
+ // Emit the init() function.
+ fmt.Fprintf(outputFile, "func init() {\n")
+ for _, ic := range initCalls {
+ fmt.Fprintf(outputFile, " %s\n", ic)
+ }
+ fmt.Fprintf(outputFile, "}\n")
+}