diff options
Diffstat (limited to 'tools/go_stateify')
-rw-r--r-- | tools/go_stateify/BUILD | 9 | ||||
-rw-r--r-- | tools/go_stateify/defs.bzl | 77 | ||||
-rw-r--r-- | tools/go_stateify/main.go | 386 |
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") +} |