Drop runtime use of gogo packages

staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/cmd.go

@@ -46,6 +46,7 @@ type Generator struct {
 	Clean                bool
 	OnlyIDL              bool
 	KeepGogoproto        bool
+	DropGogoGo           bool
 	SkipGeneratedRewrite bool
 	DropEmbeddedFields   string
 }
@@ -65,6 +66,7 @@ func New() *Generator {
 		}, ","),
 		Packages:           "",
 		DropEmbeddedFields: "k8s.io/apimachinery/pkg/apis/meta/v1.TypeMeta",
+		DropGogoGo:         true,
 	}
 }
 
@@ -79,6 +81,7 @@ func (g *Generator) BindFlags(flag *flag.FlagSet) {
 	flag.BoolVar(&g.OnlyIDL, "only-idl", g.OnlyIDL, "If true, only generate the IDL for each package.")
 	flag.BoolVar(&g.KeepGogoproto, "keep-gogoproto", g.KeepGogoproto, "If true, the generated IDL will contain gogoprotobuf extensions which are normally removed")
 	flag.BoolVar(&g.SkipGeneratedRewrite, "skip-generated-rewrite", g.SkipGeneratedRewrite, "If true, skip fixing up the generated.pb.go file (debugging only).")
+	flag.BoolVar(&g.DropGogoGo, "drop-gogo-go", g.DropGogoGo, "Drop all references to gogo packages in generated code")
 	flag.StringVar(&g.DropEmbeddedFields, "drop-embedded-fields", g.DropEmbeddedFields, "Comma-delimited list of embedded Go types to omit from generated protobufs")
 }
 
@@ -99,6 +102,10 @@ func Run(g *Generator) {
 		log.Fatalf("Both apimachinery-packages and packages are empty. At least one package must be specified.")
 	}
 
+	if g.DropGogoGo && g.SkipGeneratedRewrite {
+		log.Fatalf("--drop-gogo-go=true and --skip-generated-rewrite=true are mutually exclusive")
+	}
+
 	// Build up a list of packages to load from all the inputs.  Track the
 	// special modifiers for each.  NOTE: This does not support pkg/... syntax.
 	type modifier struct {
@@ -288,7 +295,7 @@ func Run(g *Generator) {
 
 		// alter the generated protobuf file to remove the generated types (but leave the serializers) and rewrite the
 		// package statement to match the desired package name
-		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, p.OptionalTypeName, buf.Bytes()); err != nil {
+		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, p.OptionalTypeName, buf.Bytes(), g.DropGogoGo); err != nil {
 			log.Fatalf("Unable to rewrite generated %s: %v", outputPath, err)
 		}
 

staging/src/k8s.io/code-generator/cmd/go-to-protobuf/protobuf/parser.go

@@ -77,7 +77,7 @@ type ExtractFunc func(*ast.TypeSpec) bool
 // and should have its marshal functions adjusted to remove the 'Items' accessor.
 type OptionalFunc func(name string) bool
 
-func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, optionalFn OptionalFunc, header []byte) error {
+func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, optionalFn OptionalFunc, header []byte, dropGogo bool) error {
 	return rewriteFile(name, header, func(fset *token.FileSet, file *ast.File) error {
 		cmap := ast.NewCommentMap(fset, file, file.Comments)
 
@@ -85,6 +85,22 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option
 		for _, d := range file.Decls {
 			rewriteOptionalMethods(d, optionalFn)
 		}
+		if dropGogo {
+			// transform references to gogo sort util
+			var oldSortImport string
+			var usedSort bool
+			for _, d := range file.Decls {
+				oldSortImport, usedSort = rewriteGogoSortImport(d)
+				if usedSort {
+					break
+				}
+			}
+			if usedSort {
+				for _, d := range file.Decls {
+					rewriteGogoSort(d, oldSortImport, "sort")
+				}
+			}
+		}
 
 		// remove types that are already declared
 		decls := []ast.Decl{}
@@ -95,6 +111,10 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option
 			if dropEmptyImportDeclarations(d) {
 				continue
 			}
+			// remove all but required functions
+			if dropGogo && dropUnusedGo(d) {
+				continue
+			}
 			decls = append(decls, d)
 		}
 		file.Decls = decls
@@ -105,6 +125,133 @@ func RewriteGeneratedGogoProtobufFile(name string, extractFn ExtractFunc, option
 	})
 }
 
+// rewriteGogoSortImport rewrites an import of "github.com/gogo/protobuf/sortkeys" to "sort",
+// and returns the original package alias and true if the rewrite occurred.
+// Returns "", false if the decl is not an import decl, or does not contain an import of "github.com/gogo/protobuf/sortkeys",
+func rewriteGogoSortImport(decl ast.Decl) (string, bool) {
+	t, ok := decl.(*ast.GenDecl)
+	if !ok {
+		return "", false
+	}
+	if t.Tok != token.IMPORT {
+		return "", false
+	}
+	for _, s := range t.Specs {
+		if spec, ok := s.(*ast.ImportSpec); ok {
+			if spec.Path != nil && spec.Path.Value == `"github.com/gogo/protobuf/sortkeys"` {
+				// switch gogo sort to stdlib sort
+				spec.Path.Value = `"sort"`
+				oldName := "sortkeys"
+				if spec.Name != nil {
+					oldName = spec.Name.Name
+				}
+				spec.Name = nil
+				return oldName, true
+			}
+		}
+	}
+	return "", false
+}
+
+// rewriteGogoSort walks the AST, replacing use of the oldSortImport package with newSortImport
+func rewriteGogoSort(decl ast.Decl, oldSortImport, newSortImport string) {
+	t, ok := decl.(*ast.FuncDecl)
+	if !ok {
+		return
+	}
+	ast.Walk(replacePackageVisitor{oldPackage: oldSortImport, newPackage: newSortImport}, t.Body)
+}
+
+// keepFuncs is an allowlist of top-level func decls we should keep
+var keepFuncs = map[string]bool{
+	// generated helpers
+	"sovGenerated":           true,
+	"sozGenerated":           true,
+	"skipGenerated":          true,
+	"encodeVarintGenerated":  true,
+	"valueToStringGenerated": true,
+
+	// unmarshal
+	"Reset":        true,
+	"ProtoMessage": true,
+	"Unmarshal":    true,
+
+	// marshal
+	"Size":                 true,
+	"Marshal":              true,
+	"MarshalTo":            true,
+	"MarshalToSizedBuffer": true,
+
+	// other widely used methods
+	"String": true,
+}
+
+// keepVars is an allowlist of top-level var decls we should keep
+var keepVars = map[string]bool{
+	"ErrInvalidLengthGenerated":        true,
+	"ErrIntOverflowGenerated":          true,
+	"ErrUnexpectedEndOfGroupGenerated": true,
+}
+
+// dropUnusedGo returns true if the top-level decl should be dropped.
+// Has the following behavior for different decl types:
+// * import: decl is rewritten to drop gogo package imports. Returns true if all imports in the decl were gogo imports, false if non-gogo imports remain.
+// * var: decl is rewritten to drop vars not in the keepVars allowlist. Returns true if all vars in the decl were removed, false if allowlisted vars remain.
+// * const: returns true
+// * type: returns true
+// * func: returns true if the func is not in the keepFuncs allowlist and should be dropped.
+// * other: returns false
+func dropUnusedGo(decl ast.Decl) bool {
+	switch t := decl.(type) {
+	case *ast.GenDecl:
+		switch t.Tok {
+		case token.IMPORT:
+			specs := []ast.Spec{}
+			for _, s := range t.Specs {
+				if spec, ok := s.(*ast.ImportSpec); ok {
+					if spec.Path == nil || !strings.HasPrefix(spec.Path.Value, `"github.com/gogo/protobuf/`) {
+						specs = append(specs, spec)
+					}
+				}
+			}
+			if len(specs) == 0 {
+				return true
+			}
+			t.Specs = specs
+			return false
+		case token.CONST:
+			// drop all const declarations
+			return true
+		case token.VAR:
+			specs := []ast.Spec{}
+			for _, s := range t.Specs {
+				if spec, ok := s.(*ast.ValueSpec); ok {
+					if keepVars[spec.Names[0].Name] {
+						specs = append(specs, spec)
+					}
+				}
+			}
+			if len(specs) == 0 {
+				return true
+			}
+			t.Specs = specs
+			return false
+		case token.TYPE:
+			// drop all type declarations
+			return true
+		}
+	case *ast.FuncDecl:
+		name := ""
+		if t.Name != nil {
+			name = t.Name.Name
+		}
+		return !keepFuncs[name]
+	default:
+		return false
+	}
+	return false
+}
+
 // rewriteOptionalMethods makes specific mutations to marshaller methods that belong to types identified
 // as being "optional" (they may be nil on the wire). This allows protobuf to serialize a map or slice and
 // properly discriminate between empty and nil (which is not possible in protobuf).
@@ -451,3 +598,19 @@ func updateStructTags(decl ast.Decl, structTags map[string]map[string]string, to
 	}
 	return errs
 }
+
+type replacePackageVisitor struct {
+	oldPackage string
+	newPackage string
+}
+
+// Visit walks the provided node, transforming references to the old package to the new package.
+func (v replacePackageVisitor) Visit(n ast.Node) ast.Visitor {
+	if e, ok := n.(*ast.SelectorExpr); ok {
+		if i, ok := e.X.(*ast.Ident); ok && i.Name == v.oldPackage {
+			i.Name = v.newPackage
+		}
+		return nil
+	}
+	return v
+}