Refactor + fixes

pkg/applyconfigurations/gen.go

@@ -24,7 +24,6 @@ import (
 	"go/types"
 	"io"
 	"path/filepath"
-	"regexp"
 	"sort"
 	"strings"
 
@@ -37,11 +36,19 @@ import (
 // Based on deepcopy gen but with legacy marker support removed.
 
 var (
-	enablePkgMarker  = markers.Must(markers.MakeDefinition("kubebuilder:object:generate", markers.DescribesPackage, false))
-	enableTypeMarker = markers.Must(markers.MakeDefinition("kubebuilder:object:generate", markers.DescribesType, false))
-	isObjectMarker   = markers.Must(markers.MakeDefinition("kubebuilder:object:root", markers.DescribesType, false))
+	enablePkgMarker  = markers.Must(markers.MakeDefinition("kubebuilder:ac:generate", markers.DescribesPackage, false))
+	enableTypeMarker = markers.Must(markers.MakeDefinition("kubebuilder:ac:generate", markers.DescribesType, false))
+	isObjectMarker   = markers.Must(markers.MakeDefinition("kubebuilder:ac:root", markers.DescribesType, false))
 )
 
+var importMapping = map[string]string{
+	"k8s.io/apimachinery/pkg/apis/": "k8s.io/client-go/applyconfigurations/",
+	"k8s.io/api/": "k8s.io/client-go/applyconfigurations/",
+}
+
+const importPathSuffix = "ac"
+const packageFileName = "zz_generated.applyconfigurations.go"
+
 // +controllertools:marker:generateHelp
 
 // Generator generates code containing apply configuration type implementations.
@@ -102,21 +109,23 @@ func genObjectInterface(info *markers.TypeInfo) bool {
 	return false
 }
 
-func groupAndPackageVersion(pkg string) string {
-	parts := strings.Split(pkg, "/")
-	return parts[len(parts)-2] + "/" + parts[len(parts)-1]
+func createApplyConfigPackage(universe *Universe, pkg *loader.Package) *loader.Package {
+	newPkg := &loader.Package{Package: &packages.Package{}}
+	dir := filepath.Dir(pkg.CompiledGoFiles[0])
+	// TODO|jefftree: This forces the package to live in a new "ac" directory. Is this code the best way to accomplish the task?
+	newPkg.CompiledGoFiles = append(newPkg.CompiledGoFiles, dir + "/" + importPathSuffix+"/")
+	return newPkg
 }
 
-func createApplyConfigPackage(universe Universe, pkg *loader.Package) *loader.Package {
-	newPkg := &loader.Package{Package: &packages.Package{}}
+type PkgInfo struct {
+	objGenCtx *ObjectGenCtx
+	pkg       *loader.Package
+	used      bool
+	typeInfo  []types.Type
+}
 
-	if filepath.Dir(pkg.CompiledGoFiles[0]) == universe.baseFilePath {
-		newPkg.CompiledGoFiles = append(newPkg.CompiledGoFiles, universe.baseFilePath+"/"+universe.importPathSuffix+"/")
-	} else {
-		desiredPath := universe.baseFilePath + "/" + universe.importPathSuffix + "/" + groupAndPackageVersion(pkg.PkgPath) + "/"
-		newPkg.CompiledGoFiles = append(newPkg.CompiledGoFiles, desiredPath)
-	}
-	return newPkg
+func (p *PkgInfo) GenerateTypes() {
+	p.typeInfo = p.objGenCtx.generateEligibleTypes(p.pkg)
 }
 
 func (d Generator) Generate(ctx *genall.GenerationContext) error {
@@ -140,29 +149,17 @@ func (d Generator) Generate(ctx *genall.GenerationContext) error {
 	var pkgList []*loader.Package
 	visited := make(map[string]*loader.Package)
 
-	//TODO|jefftree: This might cause problems if multiple packages are provided
-	crdRoot := ctx.Roots[0]
-
 	for _, root := range ctx.Roots {
+		visited[root.PkgPath] = root
 		pkgList = append(pkgList, root)
 	}
 
-	for len(pkgList) != 0 {
-		pkg := pkgList[0]
-		pkgList = pkgList[1:]
-		if _, ok := visited[pkg.PkgPath]; ok {
-			continue
-		}
-
-		visited[pkg.PkgPath] = pkg
-
+	for _, pkg := range pkgList {
 		for _, imp := range pkg.Imports() {
-			// Only index k8s types
-			match, _ := regexp.MatchString("k8s.io/.*apis?/.+", imp.PkgPath)
-			if !match {
+			if _, ok := visited[imp.PkgPath]; ok {
 				continue
 			}
-			pkgList = append(pkgList, imp)
+			visited[imp.PkgPath] = imp
 		}
 	}
 
@@ -171,21 +168,15 @@ func (d Generator) Generate(ctx *genall.GenerationContext) error {
 		eligibleTypes = append(eligibleTypes, objGenCtx.generateEligibleTypes(pkg)...)
 	}
 
-	universe := Universe{
-		eligibleTypes:    eligibleTypes,
-		baseImportPath:   crdRoot.PkgPath,
-		importPathSuffix: "ac",
-		baseFilePath:     filepath.Dir(crdRoot.CompiledGoFiles[0]),
+	universe := &Universe{
+		eligibleTypes:  eligibleTypes,
 	}
 
-	// universe.baseImportPath = "k8s.io/client-go/applyconfigurations"
-
-	for _, pkg := range visited {
+	for _, pkg := range pkgList {
 		outContents := objGenCtx.generateForPackage(universe, pkg)
 		if outContents == nil {
 			continue
 		}
-
 		newPkg := createApplyConfigPackage(universe, pkg)
 		writeOut(ctx, newPkg, outContents)
 	}
@@ -242,15 +233,49 @@ func (ctx *ObjectGenCtx) generateEligibleTypes(root *loader.Package) []types.Typ
 
 type Universe struct {
 	eligibleTypes    []types.Type
-	baseImportPath   string
-	importPathSuffix string
-	baseFilePath     string
+}
+
+func (u *Universe) existingApplyConfig(typeInfo *types.Named, pkgPath string) (string, bool) {
+	for prefix, replacePath := range importMapping {
+		if strings.HasPrefix(pkgPath, prefix) {
+			path := replacePath + strings.TrimPrefix(pkgPath, prefix)
+			return path, true
+		}
+	}
+	return "", false
+}
+
+func (u *Universe) IsApplyConfigGenerated(typeInfo *types.Named, pkgPath string) bool {
+	exists := false
+	for _, b := range u.eligibleTypes {
+		if b == typeInfo {
+			exists = true
+			break
+		}
+	}
+	return exists
+}
+
+func (u *Universe) GetApplyConfigPath(typeInfo *types.Named, pkgPath string) (string, bool) {
+	isApplyConfigGenerated := u.IsApplyConfigGenerated(typeInfo, pkgPath)
+	if path, ok := u.existingApplyConfig(typeInfo, pkgPath); ok {
+		if isApplyConfigGenerated {
+			return path, true
+		} else {
+			return pkgPath, false
+		}
+	}
+	// ApplyConfig is necessary but location is not explicitly specified. Assume the ApplyConfig exists at the below directory
+	if isApplyConfigGenerated {
+		return pkgPath + "/" + importPathSuffix, true
+	}
+	return pkgPath, false
 }
 
 // generateForPackage generates apply configuration implementations for
 // types in the given package, writing the formatted result to given writer.
 // May return nil if source could not be generated.
-func (ctx *ObjectGenCtx) generateForPackage(universe Universe, root *loader.Package) []byte {
+func (ctx *ObjectGenCtx) generateForPackage(universe *Universe, root *loader.Package) []byte {
 	byType := make(map[string][]byte)
 	imports := &importsList{
 		byPath:  make(map[string]string),
@@ -274,17 +299,17 @@ func (ctx *ObjectGenCtx) generateForPackage(universe Universe, root *loader.Pack
 			codeWriter:  &codeWriter{out: outContent},
 		}
 
-		copyCtx.GenerateTypesFor(&universe, root, info)
+		copyCtx.GenerateTypesFor(universe, root, info)
 		for _, field := range info.Fields {
 			if field.Name != "" {
-				copyCtx.GenerateMemberSet(&universe, field, root, info)
+				copyCtx.GenerateMemberSet(universe, field, root, info)
 			}
 		}
 
 		copyCtx.GenerateStructConstructor(root, info)
 
 		if isRootType(info) {
-			copyCtx.GenerateTypeGetter(&universe, root, info)
+			copyCtx.GenerateRootFunctions(universe, root, info)
 		}
 
 		outBytes := outContent.Bytes()
@@ -335,7 +360,7 @@ func writeTypes(pkg *loader.Package, out io.Writer, byType map[string][]byte) {
 // writeFormatted outputs the given code, after gofmt-ing it.  If we couldn't gofmt,
 // we write the unformatted code for debugging purposes.
 func writeOut(ctx *genall.GenerationContext, root *loader.Package, outBytes []byte) {
-	outputFile, err := ctx.Open(root, "zz_generated.applyconfigurations.go")
+	outputFile, err := ctx.Open(root, packageFileName)
 	if err != nil {
 		root.AddError(err)
 		return