Skip to content

Commit

Permalink
stateify: Add +stateify identtype tag for identical type definiti…
Browse files Browse the repository at this point in the history
…ons.

This allows serialization of type definitions identical of stateifiable
struct types.

Context: I want to add a `bpf.Instruction` library which is a type definition
of `linux.BPFInstruction` (but with nicer helper functions), and keep its
serializability.

Sample output for

```go
// +stateify savable
// +stateify alias:linux.BPFInstruction
type Instruction linux.BPFInstruction
```

```go
func (ins *Instruction) StateTypeName() string {
        return "pkg/bpf.Instruction"
}

func (ins *Instruction) StateFields() []string {
        return (*linux.BPFInstruction)(ins).StateFields()
}

// +checklocksignore
func (ins *Instruction) StateSave(stateSinkObject state.Sink) {
        (*linux.BPFInstruction)(ins).StateSave(stateSinkObject)
}

// +checklocksignore
func (ins *Instruction) StateLoad(stateSourceObject state.Source) {
        (*linux.BPFInstruction)(ins).StateLoad(stateSourceObject)
}
```

PiperOrigin-RevId: 568999266
  • Loading branch information
EtiennePerot authored and gvisor-bot committed Sep 27, 2023
1 parent 7594667 commit 0e18a82
Showing 1 changed file with 47 additions and 4 deletions.
51 changes: 47 additions & 4 deletions tools/go_stateify/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -314,22 +314,29 @@ func main() {
// savable" in one of the proceeding comment lines. If
// the line is marked "// +stateify type" then only
// generate type information and register the type.
// If the type also has a "// +stateify identtype"
// comment, the functions are instead generated to refer to
// the type that this newly-defined type is identical to, rather
// than about the newly-defined type itself.
if d.Doc == nil {
continue
}
var (
generateTypeInfo = false
generateSaverLoader = false
isIdentType = false
)
for _, l := range d.Doc.List {
if l.Text == "// +stateify savable" {
generateTypeInfo = true
generateSaverLoader = true
break
}
if l.Text == "// +stateify type" {
generateTypeInfo = true
}
if l.Text == "// +stateify identtype" {
isIdentType = true
}
}
if !generateTypeInfo && !generateSaverLoader {
continue
Expand All @@ -345,6 +352,10 @@ func main() {
switch x := ts.Type.(type) {
case *ast.StructType:
maybeEmitImports()
if isIdentType {
fmt.Fprintf(os.Stderr, "Cannot use `+stateify identtype` on a struct type (%v); must be a type definition of an identical type.", ts.Name.Name)
os.Exit(1)
}

// Record the slot for each field.
fieldCount := 0
Expand Down Expand Up @@ -463,9 +474,41 @@ func main() {
fmt.Fprintf(outputFile, "func (%s *%s) StateTypeName() string {\n", recv, ts.Name.Name)
fmt.Fprintf(outputFile, " return \"%s.%s\"\n", *fullPkg, ts.Name.Name)
fmt.Fprintf(outputFile, "}\n\n")
fmt.Fprintf(outputFile, "func (%s *%s) StateFields() []string {\n", recv, ts.Name.Name)
fmt.Fprintf(outputFile, " return nil\n")
fmt.Fprintf(outputFile, "}\n\n")

if !isIdentType {
fmt.Fprintf(outputFile, "func (%s *%s) StateFields() []string {\n", recv, ts.Name.Name)
fmt.Fprintf(outputFile, " return nil\n")
fmt.Fprintf(outputFile, "}\n\n")
} else {
var typeName string
switch y := x.(type) {
case *ast.Ident:
typeName = y.Name
case *ast.SelectorExpr:
expIdent, ok := y.X.(*ast.Ident)
if !ok {
fmt.Fprintf(os.Stderr, "Cannot use non-ident %v (type %T) in type selector expression %v", y.X, y.X, y)
os.Exit(1)
}
typeName = fmt.Sprintf("%s.%s", expIdent.Name, y.Sel.Name)
default:
fmt.Fprintf(os.Stderr, "Cannot use `+stateify identtype` on a non-identifier/non-selector type definition (%v => %v of type %T); must be a type definition of an identical type.", ts.Name.Name, x, x)
os.Exit(1)
}
fmt.Fprintf(outputFile, "func (%s *%s) StateFields() []string {\n", recv, ts.Name.Name)
fmt.Fprintf(outputFile, " return (*%s)(%s).StateFields()\n", typeName, recv)
fmt.Fprintf(outputFile, "}\n\n")
if generateSaverLoader {
fmt.Fprintf(outputFile, "// +checklocksignore\n")
fmt.Fprintf(outputFile, "func (%s *%s) StateSave(stateSinkObject %sSink) {\n", recv, ts.Name.Name, statePrefix)
fmt.Fprintf(outputFile, " (*%s)(%s).StateSave(stateSinkObject)\n", typeName, recv)
fmt.Fprintf(outputFile, "}\n\n")
fmt.Fprintf(outputFile, "// +checklocksignore\n")
fmt.Fprintf(outputFile, "func (%s *%s) StateLoad(stateSourceObject %sSource) {\n", recv, ts.Name.Name, statePrefix)
fmt.Fprintf(outputFile, " (*%s)(%s).StateLoad(stateSourceObject)\n", typeName, recv)
fmt.Fprintf(outputFile, "}\n\n")
}
}

// See above.
emitRegister(ts.Name.Name)
Expand Down

0 comments on commit 0e18a82

Please sign in to comment.