From 0e18a826d91d20409ef32cd98c9b0512aadeb93d Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Wed, 27 Sep 2023 16:53:45 -0700
Subject: [PATCH] `stateify`: Add `+stateify identtype` tag for identical type
 definitions.

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
---
 tools/go_stateify/main.go | 51 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 4 deletions(-)

diff --git a/tools/go_stateify/main.go b/tools/go_stateify/main.go
index 4f2e39f6b..b07b76e06 100644
--- a/tools/go_stateify/main.go
+++ b/tools/go_stateify/main.go
@@ -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
@@ -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
@@ -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)