[go-marshal] Add dynamic tag in go_marshal.

This makes it easier to implement dynamically sized types in go-marshal. You
really only need to implement MarshalBytes, UnmarshalBytes and SizeBytes to
implement the entire interface.

By using the `dynamic` tag, the autogenerator will generate the rest of the
methods for us.

This change also simplifies how KernelIPTGetEntries implements Marshallable
using the newly added utility.

PiperOrigin-RevId: 356397114

3 files changed, 212 insertions, 170 deletions

diff --git a/tools/go_marshal/gomarshal/generator.go b/tools/go_marshal/gomarshal/generator.go
index fa642c88a..abd6f69ea 100644
--- a/tools/go_marshal/gomarshal/generator.go
+++ b/tools/go_marshal/gomarshal/generator.go
@@ -211,9 +211,10 @@ type sliceAPI struct {
 // marshallableType carries information about a type marked with the '+marshal'
 // directive.
 type marshallableType struct {
-	spec  *ast.TypeSpec
-	slice *sliceAPI
-	recv  string
+	spec    *ast.TypeSpec
+	slice   *sliceAPI
+	recv    string
+	dynamic bool
 }
 
 func newMarshallableType(fset *token.FileSet, tagLine *ast.Comment, spec *ast.TypeSpec) *marshallableType {
@@ -248,6 +249,9 @@ func newMarshallableType(fset *token.FileSet, tagLine *ast.Comment, spec *ast.Ty
 			}
 
 			continue
+		} else if tag == "dynamic" {
+			mt.dynamic = true
+			continue
 		}
 
 		unhandledTags = append(unhandledTags, tag)
@@ -379,23 +383,38 @@ func (g *Generator) generateOne(t *marshallableType, fset *token.FileSet) *inter
 	i := newInterfaceGenerator(t.spec, t.recv, fset)
 	switch ty := t.spec.Type.(type) {
 	case *ast.StructType:
+		if t.dynamic {
+			// Don't validate because this type is dynamically sized and probably
+			// contains some funky slices which the validation does not allow.
+			i.emitMarshallableForStruct(ty, t.dynamic)
+			if t.slice != nil {
+				abortAt(fset.Position(t.slice.comment.Slash), "Slice API is not supported for dynamic types because it assumes that each slice element is statically sized.")
+			}
+			break
+		}
 		i.validateStruct(t.spec, ty)
-		i.emitMarshallableForStruct(ty)
+		i.emitMarshallableForStruct(ty, t.dynamic)
 		if t.slice != nil {
 			i.emitMarshallableSliceForStruct(ty, t.slice)
 		}
 	case *ast.Ident:
 		i.validatePrimitiveNewtype(ty)
+		if t.dynamic {
+			abortAt(fset.Position(t.slice.comment.Slash), "Primitive type marked as '+marshal dynamic', but primitive types can not be dynamic.")
+		}
 		i.emitMarshallableForPrimitiveNewtype(ty)
 		if t.slice != nil {
 			i.emitMarshallableSliceForPrimitiveNewtype(ty, t.slice)
 		}
 	case *ast.ArrayType:
 		i.validateArrayNewtype(t.spec.Name, ty)
+		if t.dynamic {
+			abortAt(fset.Position(t.slice.comment.Slash), "Marking array types as `dynamic` is currently not supported.")
+		}
 		// After validate, we can safely call arrayLen.
 		i.emitMarshallableForArrayNewtype(t.spec.Name, ty, ty.Elt.(*ast.Ident))
 		if t.slice != nil {
-			abortAt(fset.Position(t.slice.comment.Slash), fmt.Sprintf("Array type marked as '+marshal slice:...', but this is not supported. Perhaps fold one of the dimensions?"))
+			abortAt(fset.Position(t.slice.comment.Slash), "Array type marked as '+marshal slice:...', but this is not supported. Perhaps fold one of the dimensions?")
 		}
 	default:
 		// This should've been filtered out by collectMarshallabeTypes.
@@ -408,7 +427,7 @@ func (g *Generator) generateOne(t *marshallableType, fset *token.FileSet) *inter
 // implementations type t.
 func (g *Generator) generateOneTestSuite(t *marshallableType) *testGenerator {
 	i := newTestGenerator(t.spec, t.recv)
-	i.emitTests(t.slice)
+	i.emitTests(t.slice, t.dynamic)
 	return i
 }
 
diff --git a/tools/go_marshal/gomarshal/generator_interfaces_struct.go b/tools/go_marshal/gomarshal/generator_interfaces_struct.go
index 5f6306b8f..f98e41ed7 100644
--- a/tools/go_marshal/gomarshal/generator_interfaces_struct.go
+++ b/tools/go_marshal/gomarshal/generator_interfaces_struct.go
@@ -69,7 +69,11 @@ func (g *interfaceGenerator) validateStruct(ts *ast.TypeSpec, st *ast.StructType
 	})
 }
 
-func (g *interfaceGenerator) isStructPacked(st *ast.StructType) bool {
+func (g *interfaceGenerator) isStructPacked(st *ast.StructType, isDynamic bool) bool {
+	if isDynamic {
+		// Dynamic types are not packed because a slice header might be present.
+		return false
+	}
 	packed := true
 	forEachStructField(st, func(f *ast.Field) {
 		if f.Tag != nil {
@@ -85,165 +89,17 @@ func (g *interfaceGenerator) isStructPacked(st *ast.StructType) bool {
 	return packed
 }
 
-func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
-	thisPacked := g.isStructPacked(st)
-
-	g.emit("// SizeBytes implements marshal.Marshallable.SizeBytes.\n")
-	g.emit("func (%s *%s) SizeBytes() int {\n", g.r, g.typeName())
-	g.inIndent(func() {
-		primitiveSize := 0
-		var dynamicSizeTerms []string
-
-		forEachStructField(st, fieldDispatcher{
-			primitive: func(_, t *ast.Ident) {
-				if size, dynamic := g.scalarSize(t); !dynamic {
-					primitiveSize += size
-				} else {
-					g.recordUsedMarshallable(t.Name)
-					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()", t.Name))
-				}
-			},
-			selector: func(_, tX, tSel *ast.Ident) {
-				tName := fmt.Sprintf("%s.%s", tX.Name, tSel.Name)
-				g.recordUsedImport(tX.Name)
-				g.recordUsedMarshallable(tName)
-				dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()", tName))
-			},
-			array: func(_ *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
-				lenExpr := g.arrayLenExpr(a)
-				if size, dynamic := g.scalarSize(t); !dynamic {
-					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("%d*%s", size, lenExpr))
-				} else {
-					g.recordUsedMarshallable(t.Name)
-					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()*%s", t.Name, lenExpr))
-				}
-			},
-		}.dispatch)
-		g.emit("return %d", primitiveSize)
-		if len(dynamicSizeTerms) > 0 {
-			g.incIndent()
-		}
-		{
-			for _, d := range dynamicSizeTerms {
-				g.emitNoIndent(" +\n")
-				g.emit(d)
-			}
-		}
-		if len(dynamicSizeTerms) > 0 {
-			g.decIndent()
-		}
-	})
-	g.emit("\n}\n\n")
+func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType, isDynamic bool) {
+	thisPacked := g.isStructPacked(st, isDynamic)
 
-	g.emit("// MarshalBytes implements marshal.Marshallable.MarshalBytes.\n")
-	g.emit("func (%s *%s) MarshalBytes(dst []byte) {\n", g.r, g.typeName())
-	g.inIndent(func() {
-		forEachStructField(st, fieldDispatcher{
-			primitive: func(n, t *ast.Ident) {
-				if n.Name == "_" {
-					g.emit("// Padding: dst[:sizeof(%s)] ~= %s(0)\n", t.Name, t.Name)
-					if len, dynamic := g.scalarSize(t); !dynamic {
-						g.shift("dst", len)
-					} else {
-						// We can't use shiftDynamic here because we don't have
-						// an instance of the dynamic type we can reference here
-						// (since the version in this struct is anonymous). Use
-						// a typed nil pointer to call SizeBytes() instead.
-						g.emit("dst = dst[(*%s)(nil).SizeBytes():]\n", t.Name)
-					}
-					return
-				}
-				g.marshalScalar(g.fieldAccessor(n), t.Name, "dst")
-			},
-			selector: func(n, tX, tSel *ast.Ident) {
-				if n.Name == "_" {
-					g.emit("// Padding: dst[:sizeof(%s)] ~= %s(0)\n", tX.Name, tSel.Name)
-					g.emit("dst = dst[(*%s.%s)(nil).SizeBytes():]\n", tX.Name, tSel.Name)
-					return
-				}
-				g.marshalScalar(g.fieldAccessor(n), fmt.Sprintf("%s.%s", tX.Name, tSel.Name), "dst")
-			},
-			array: func(n *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
-				lenExpr := g.arrayLenExpr(a)
-				if n.Name == "_" {
-					g.emit("// Padding: dst[:sizeof(%s)*%s] ~= [%s]%s{0}\n", t.Name, lenExpr, lenExpr, t.Name)
-					if size, dynamic := g.scalarSize(t); !dynamic {
-						g.emit("dst = dst[%d*(%s):]\n", size, lenExpr)
-					} else {
-						// We can't use shiftDynamic here because we don't have
-						// an instance of the dynamic type we can reference here
-						// (since the version in this struct is anonymous). Use
-						// a typed nil pointer to call SizeBytes() instead.
-						g.emit("dst = dst[(*%s)(nil).SizeBytes()*(%s):]\n", t.Name, lenExpr)
-					}
-					return
-				}
-
-				g.emit("for idx := 0; idx < %s; idx++ {\n", lenExpr)
-				g.inIndent(func() {
-					g.marshalScalar(fmt.Sprintf("%s[idx]", g.fieldAccessor(n)), t.Name, "dst")
-				})
-				g.emit("}\n")
-			},
-		}.dispatch)
-	})
-	g.emit("}\n\n")
-
-	g.emit("// UnmarshalBytes implements marshal.Marshallable.UnmarshalBytes.\n")
-	g.emit("func (%s *%s) UnmarshalBytes(src []byte) {\n", g.r, g.typeName())
-	g.inIndent(func() {
-		forEachStructField(st, fieldDispatcher{
-			primitive: func(n, t *ast.Ident) {
-				if n.Name == "_" {
-					g.emit("// Padding: var _ %s ~= src[:sizeof(%s)]\n", t.Name, t.Name)
-					if len, dynamic := g.scalarSize(t); !dynamic {
-						g.shift("src", len)
-					} else {
-						// We don't have an instance of the dynamic type we can
-						// reference here (since the version in this struct is
-						// anonymous). Use a typed nil pointer to call
-						// SizeBytes() instead.
-						g.shiftDynamic("src", fmt.Sprintf("(*%s)(nil)", t.Name))
-						g.recordPotentiallyNonPackedField(fmt.Sprintf("(*%s)(nil)", t.Name))
-					}
-					return
-				}
-				g.unmarshalScalar(g.fieldAccessor(n), t.Name, "src")
-			},
-			selector: func(n, tX, tSel *ast.Ident) {
-				if n.Name == "_" {
-					g.emit("// Padding: %s ~= src[:sizeof(%s.%s)]\n", g.fieldAccessor(n), tX.Name, tSel.Name)
-					g.emit("src = src[(*%s.%s)(nil).SizeBytes():]\n", tX.Name, tSel.Name)
-					g.recordPotentiallyNonPackedField(fmt.Sprintf("(*%s.%s)(nil)", tX.Name, tSel.Name))
-					return
-				}
-				g.unmarshalScalar(g.fieldAccessor(n), fmt.Sprintf("%s.%s", tX.Name, tSel.Name), "src")
-			},
-			array: func(n *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
-				lenExpr := g.arrayLenExpr(a)
-				if n.Name == "_" {
-					g.emit("// Padding: ~ copy([%s]%s(%s), src[:sizeof(%s)*%s])\n", lenExpr, t.Name, g.fieldAccessor(n), t.Name, lenExpr)
-					if size, dynamic := g.scalarSize(t); !dynamic {
-						g.emit("src = src[%d*(%s):]\n", size, lenExpr)
-					} else {
-						// We can't use shiftDynamic here because we don't have
-						// an instance of the dynamic type we can referece here
-						// (since the version in this struct is anonymous). Use
-						// a typed nil pointer to call SizeBytes() instead.
-						g.emit("src = src[(*%s)(nil).SizeBytes()*(%s):]\n", t.Name, lenExpr)
-					}
-					return
-				}
-
-				g.emit("for idx := 0; idx < %s; idx++ {\n", lenExpr)
-				g.inIndent(func() {
-					g.unmarshalScalar(fmt.Sprintf("%s[idx]", g.fieldAccessor(n)), t.Name, "src")
-				})
-				g.emit("}\n")
-			},
-		}.dispatch)
-	})
-	g.emit("}\n\n")
+	// Dynamic types are supposed to manually implement SizeBytes, MarshalBytes
+	// and UnmarshalBytes. The rest of the methos are autogenerated and depend on
+	// the implementation of these three.
+	if !isDynamic {
+		g.emitSizeBytesForStruct(st)
+		g.emitMarshalBytesForStruct(st)
+		g.emitUnmarshalBytesForStruct(st)
+	}
 
 	g.emit("// Packed implements marshal.Marshallable.Packed.\n")
 	g.emit("//go:nosplit\n")
@@ -428,8 +284,171 @@ func (g *interfaceGenerator) emitMarshallableForStruct(st *ast.StructType) {
 	g.emit("}\n\n")
 }
 
+func (g *interfaceGenerator) emitSizeBytesForStruct(st *ast.StructType) {
+	g.emit("// SizeBytes implements marshal.Marshallable.SizeBytes.\n")
+	g.emit("func (%s *%s) SizeBytes() int {\n", g.r, g.typeName())
+	g.inIndent(func() {
+		primitiveSize := 0
+		var dynamicSizeTerms []string
+
+		forEachStructField(st, fieldDispatcher{
+			primitive: func(_, t *ast.Ident) {
+				if size, dynamic := g.scalarSize(t); !dynamic {
+					primitiveSize += size
+				} else {
+					g.recordUsedMarshallable(t.Name)
+					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()", t.Name))
+				}
+			},
+			selector: func(_, tX, tSel *ast.Ident) {
+				tName := fmt.Sprintf("%s.%s", tX.Name, tSel.Name)
+				g.recordUsedImport(tX.Name)
+				g.recordUsedMarshallable(tName)
+				dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()", tName))
+			},
+			array: func(_ *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
+				lenExpr := g.arrayLenExpr(a)
+				if size, dynamic := g.scalarSize(t); !dynamic {
+					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("%d*%s", size, lenExpr))
+				} else {
+					g.recordUsedMarshallable(t.Name)
+					dynamicSizeTerms = append(dynamicSizeTerms, fmt.Sprintf("(*%s)(nil).SizeBytes()*%s", t.Name, lenExpr))
+				}
+			},
+		}.dispatch)
+		g.emit("return %d", primitiveSize)
+		if len(dynamicSizeTerms) > 0 {
+			g.incIndent()
+		}
+		{
+			for _, d := range dynamicSizeTerms {
+				g.emitNoIndent(" +\n")
+				g.emit(d)
+			}
+		}
+		if len(dynamicSizeTerms) > 0 {
+			g.decIndent()
+		}
+	})
+	g.emit("\n}\n\n")
+}
+
+func (g *interfaceGenerator) emitMarshalBytesForStruct(st *ast.StructType) {
+	g.emit("// MarshalBytes implements marshal.Marshallable.MarshalBytes.\n")
+	g.emit("func (%s *%s) MarshalBytes(dst []byte) {\n", g.r, g.typeName())
+	g.inIndent(func() {
+		forEachStructField(st, fieldDispatcher{
+			primitive: func(n, t *ast.Ident) {
+				if n.Name == "_" {
+					g.emit("// Padding: dst[:sizeof(%s)] ~= %s(0)\n", t.Name, t.Name)
+					if len, dynamic := g.scalarSize(t); !dynamic {
+						g.shift("dst", len)
+					} else {
+						// We can't use shiftDynamic here because we don't have
+						// an instance of the dynamic type we can reference here
+						// (since the version in this struct is anonymous). Use
+						// a typed nil pointer to call SizeBytes() instead.
+						g.emit("dst = dst[(*%s)(nil).SizeBytes():]\n", t.Name)
+					}
+					return
+				}
+				g.marshalScalar(g.fieldAccessor(n), t.Name, "dst")
+			},
+			selector: func(n, tX, tSel *ast.Ident) {
+				if n.Name == "_" {
+					g.emit("// Padding: dst[:sizeof(%s)] ~= %s(0)\n", tX.Name, tSel.Name)
+					g.emit("dst = dst[(*%s.%s)(nil).SizeBytes():]\n", tX.Name, tSel.Name)
+					return
+				}
+				g.marshalScalar(g.fieldAccessor(n), fmt.Sprintf("%s.%s", tX.Name, tSel.Name), "dst")
+			},
+			array: func(n *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
+				lenExpr := g.arrayLenExpr(a)
+				if n.Name == "_" {
+					g.emit("// Padding: dst[:sizeof(%s)*%s] ~= [%s]%s{0}\n", t.Name, lenExpr, lenExpr, t.Name)
+					if size, dynamic := g.scalarSize(t); !dynamic {
+						g.emit("dst = dst[%d*(%s):]\n", size, lenExpr)
+					} else {
+						// We can't use shiftDynamic here because we don't have
+						// an instance of the dynamic type we can reference here
+						// (since the version in this struct is anonymous). Use
+						// a typed nil pointer to call SizeBytes() instead.
+						g.emit("dst = dst[(*%s)(nil).SizeBytes()*(%s):]\n", t.Name, lenExpr)
+					}
+					return
+				}
+
+				g.emit("for idx := 0; idx < %s; idx++ {\n", lenExpr)
+				g.inIndent(func() {
+					g.marshalScalar(fmt.Sprintf("%s[idx]", g.fieldAccessor(n)), t.Name, "dst")
+				})
+				g.emit("}\n")
+			},
+		}.dispatch)
+	})
+	g.emit("}\n\n")
+}
+
+func (g *interfaceGenerator) emitUnmarshalBytesForStruct(st *ast.StructType) {
+	g.emit("// UnmarshalBytes implements marshal.Marshallable.UnmarshalBytes.\n")
+	g.emit("func (%s *%s) UnmarshalBytes(src []byte) {\n", g.r, g.typeName())
+	g.inIndent(func() {
+		forEachStructField(st, fieldDispatcher{
+			primitive: func(n, t *ast.Ident) {
+				if n.Name == "_" {
+					g.emit("// Padding: var _ %s ~= src[:sizeof(%s)]\n", t.Name, t.Name)
+					if len, dynamic := g.scalarSize(t); !dynamic {
+						g.shift("src", len)
+					} else {
+						// We don't have an instance of the dynamic type we can
+						// reference here (since the version in this struct is
+						// anonymous). Use a typed nil pointer to call
+						// SizeBytes() instead.
+						g.shiftDynamic("src", fmt.Sprintf("(*%s)(nil)", t.Name))
+						g.recordPotentiallyNonPackedField(fmt.Sprintf("(*%s)(nil)", t.Name))
+					}
+					return
+				}
+				g.unmarshalScalar(g.fieldAccessor(n), t.Name, "src")
+			},
+			selector: func(n, tX, tSel *ast.Ident) {
+				if n.Name == "_" {
+					g.emit("// Padding: %s ~= src[:sizeof(%s.%s)]\n", g.fieldAccessor(n), tX.Name, tSel.Name)
+					g.emit("src = src[(*%s.%s)(nil).SizeBytes():]\n", tX.Name, tSel.Name)
+					g.recordPotentiallyNonPackedField(fmt.Sprintf("(*%s.%s)(nil)", tX.Name, tSel.Name))
+					return
+				}
+				g.unmarshalScalar(g.fieldAccessor(n), fmt.Sprintf("%s.%s", tX.Name, tSel.Name), "src")
+			},
+			array: func(n *ast.Ident, a *ast.ArrayType, t *ast.Ident) {
+				lenExpr := g.arrayLenExpr(a)
+				if n.Name == "_" {
+					g.emit("// Padding: ~ copy([%s]%s(%s), src[:sizeof(%s)*%s])\n", lenExpr, t.Name, g.fieldAccessor(n), t.Name, lenExpr)
+					if size, dynamic := g.scalarSize(t); !dynamic {
+						g.emit("src = src[%d*(%s):]\n", size, lenExpr)
+					} else {
+						// We can't use shiftDynamic here because we don't have
+						// an instance of the dynamic type we can referece here
+						// (since the version in this struct is anonymous). Use
+						// a typed nil pointer to call SizeBytes() instead.
+						g.emit("src = src[(*%s)(nil).SizeBytes()*(%s):]\n", t.Name, lenExpr)
+					}
+					return
+				}
+
+				g.emit("for idx := 0; idx < %s; idx++ {\n", lenExpr)
+				g.inIndent(func() {
+					g.unmarshalScalar(fmt.Sprintf("%s[idx]", g.fieldAccessor(n)), t.Name, "src")
+				})
+				g.emit("}\n")
+			},
+		}.dispatch)
+	})
+	g.emit("}\n\n")
+}
+
 func (g *interfaceGenerator) emitMarshallableSliceForStruct(st *ast.StructType, slice *sliceAPI) {
-	thisPacked := g.isStructPacked(st)
+	thisPacked := g.isStructPacked(st, false /* isDynamic */)
 
 	if slice.inner {
 		abortAt(g.f.Position(slice.comment.Slash), fmt.Sprintf("The ':inner' argument to '+marshal slice:%s:inner' is only applicable to newtypes on primitives. Remove it from this struct declaration.", slice.ident))
diff --git a/tools/go_marshal/gomarshal/generator_tests.go b/tools/go_marshal/gomarshal/generator_tests.go
index 6cf00843f..ca3e15c16 100644
--- a/tools/go_marshal/gomarshal/generator_tests.go
+++ b/tools/go_marshal/gomarshal/generator_tests.go
@@ -216,12 +216,16 @@ func (g *testGenerator) emitTestSizeBytesOnTypedNilPtr() {
 	})
 }
 
-func (g *testGenerator) emitTests(slice *sliceAPI) {
+func (g *testGenerator) emitTests(slice *sliceAPI, isDynamic bool) {
 	g.emitTestNonZeroSize()
 	g.emitTestSuspectAlignment()
-	g.emitTestMarshalUnmarshalPreservesData()
-	g.emitTestWriteToUnmarshalPreservesData()
-	g.emitTestSizeBytesOnTypedNilPtr()
+	if !isDynamic {
+		// Do not test these for dynamic structs because they violate some
+		// assumptions that these tests make.
+		g.emitTestMarshalUnmarshalPreservesData()
+		g.emitTestWriteToUnmarshalPreservesData()
+		g.emitTestSizeBytesOnTypedNilPtr()
+	}
 
 	if slice != nil {
 		g.emitTestMarshalUnmarshalSlicePreservesData(slice)