Omit types that aren't referred to from a resource (#170)

* Add StripUnusedDefinitions pass to CodeGenerator.Generate

Update expected output to remove unreferenced types in the expected
output for AllOf_generates_wrapper_type

* Add Referees method to Type and impls

The name is awkward - it might be better to replace References with
this.

* Implement collecting referrers (I think)

* Implement the connectedness check

If a definition isn't connected by a series of references to any root
definition (ie, a resource) then it should be filtered out.

Memoised to save checking a type name twice.

* Correct implementation of TypeName.Referees

Originally implementation was based on a misunderstanding of how the
different astmodel types worked.

* Strip unused definitions in jsonast golden tests

This required moving StripUnusedDefinitions out of pkg/codegen - I put
it into pkg/astmodel because that's where all of the types it
interacts with are. Also enabled passing in a different set of roots,
since the golden tests use the Test type (which isn't a resource) as
the root.

* Update Type.References to return a TypeNameSet, remove Referees

Returning a set makes the old behaviour simple to implement in the few
places that used it.

Adds a SetUnion to combine two TypeNameSets. I thought about making
this a method to add all of the elements of another set to the
receiver - it would be more efficient but I think it's less clean.

* Don't memoise negative checkWithPath results

These can come from a search that is unwinding because it detected a
cycle, rather than one that didn't find a root. In that case we can't
know whether a search from this node with a different path would
succeed.

Added a unit test to explore this problem.

Made the TypeNameSet methods take TypeNames rather than *TypeNames -
this simplified them a bit (at the cost of having to do some extra nil
checks at the edges).

* Corrected grammar on doc comments

* Simplify connectedness-check thanks to a suggestion from Matt

Rather than walking up the references until we find a root, walk down
the references from the roots collecting up all of the types we
reach. This is much simpler and doesn't need any memoisation (since
the accumulated set replaces it).

* Update Function.References to match Type.References

* Remove unnecessary nil checks

These *TypeNames can't be nil - there are a lot of earlier points
in the code that do no-look dereferences.

hack/generator/pkg/astmodel/array_type.go

@@ -34,9 +34,9 @@ func (array *ArrayType) RequiredImports() []*PackageReference {
 	return array.element.RequiredImports()
 }
 
-// References this type has to the given type
-func (array *ArrayType) References(d *TypeName) bool {
-	return array.element.References(d)
+// References returns the references of the type this array contains.
+func (array *ArrayType) References() TypeNameSet {
+	return array.element.References()
 }
 
 // Equals returns true if the passed type is an array type with the same kind of elements, false otherwise

hack/generator/pkg/astmodel/enum_type.go

@@ -39,9 +39,9 @@ func (enum *EnumType) AsType(codeGenerationContext *CodeGenerationContext) ast.E
 	return enum.BaseType.AsType(codeGenerationContext)
 }
 
-// References indicates whether this Type includes any direct references to the given Type
-func (enum *EnumType) References(tn *TypeName) bool {
-	return enum.BaseType.References(tn)
+// References returns any types the underlying type refers to.
+func (enum *EnumType) References() TypeNameSet {
+	return enum.BaseType.References()
 }
 
 // Equals will return true if the supplied type has the same base type and options

hack/generator/pkg/astmodel/function.go

@@ -13,7 +13,8 @@ import (
 type Function interface {
 	RequiredImports() []*PackageReference
 
-	References(name *TypeName) bool
+	// References returns the set of types to which this function refers.
+	References() TypeNameSet
 
 	// AsFunc renders the current instance as a Go abstract syntax tree
 	AsFunc(codeGenerationContext *CodeGenerationContext, receiver *TypeName, methodName string) *ast.FuncDecl

hack/generator/pkg/astmodel/map_type.go

@@ -44,9 +44,9 @@ func (m *MapType) RequiredImports() []*PackageReference {
 	return result
 }
 
-// References this type has to the given type
-func (m *MapType) References(d *TypeName) bool {
-	return m.key.References(d) || m.value.References(d)
+// References returns all of the types referenced by either the the key or value types.
+func (m *MapType) References() TypeNameSet {
+	return SetUnion(m.key.References(), m.value.References())
 }
 
 // Equals returns true if the passed type is a map type with the same kinds of keys and elements, false otherwise

hack/generator/pkg/astmodel/one_of_json_marshal_function.go

@@ -35,10 +35,10 @@ func (f *OneOfJSONMarshalFunction) Equals(other Function) bool {
 	return false
 }
 
-// References indicates whether this function includes any direct references to the given type
-func (f *OneOfJSONMarshalFunction) References(name *TypeName) bool {
+// References returns the set of references for the underlying struct.
+func (f *OneOfJSONMarshalFunction) References() TypeNameSet {
 	// Defer this check to the owning struct as we only refer to its fields and it
-	return f.oneOfStruct.References(name)
+	return f.oneOfStruct.References()
 }
 
 // AsFunc returns the function as a go ast

hack/generator/pkg/astmodel/optional_type.go

@@ -39,9 +39,9 @@ func (optional *OptionalType) RequiredImports() []*PackageReference {
 	return optional.element.RequiredImports()
 }
 
-// References is true if it is this type or the 'element' type references it
-func (optional *OptionalType) References(d *TypeName) bool {
-	return optional.element.References(d)
+// References returns the set of types that the underlying type refers to directly.
+func (optional *OptionalType) References() TypeNameSet {
+	return optional.element.References()
 }
 
 // Equals returns true if this type is equal to the other type

hack/generator/pkg/astmodel/package_definition.go

@@ -86,7 +86,10 @@ func emitFiles(filesToGenerate map[string][]TypeDefiner, outputDir string) error
 
 func anyReferences(defs []TypeDefiner, defName *TypeName) bool {
 	for _, def := range defs {
-		if def.Type().References(defName) {
+		if def == nil {
+			continue
+		}
+		if def.Type().References().Contains(*defName) {
 			return true
 		}
 	}

hack/generator/pkg/astmodel/primitive_type.go

@@ -42,10 +42,10 @@ func (prim *PrimitiveType) RequiredImports() []*PackageReference {
 	return nil
 }
 
-// References this type has to the given type
-func (prim *PrimitiveType) References(d *TypeName) bool {
-	// Primitive types dont have references
-	return false
+// References always returns nil because primitive types don't refer to
+// any other types.
+func (prim *PrimitiveType) References() TypeNameSet {
+	return nil
 }
 
 // Equals returns true if the passed type is another primitive type the same name, false otherwise

hack/generator/pkg/astmodel/strip_unused.go

@@ -0,0 +1,125 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+package astmodel
+
+// TypeNameSet stores type names in no particular order without
+// duplicates.
+type TypeNameSet map[TypeName]struct{}
+
+// NewTypeNameSet makes a TypeNameSet containing the specified
+// names. If no elements are passed it might be nil.
+func NewTypeNameSet(initial ...TypeName) TypeNameSet {
+	var result TypeNameSet
+	for _, name := range initial {
+		result = result.Add(name)
+	}
+	return result
+}
+
+// Add includes the passed name in the set and returns the updated
+// set, so that adding can work for a nil set - this makes it more
+// convenient to add to sets kept in a map (in the way you might with
+// a map of slices).
+func (ts TypeNameSet) Add(val TypeName) TypeNameSet {
+	if ts == nil {
+		ts = make(TypeNameSet)
+	}
+	ts[val] = struct{}{}
+	return ts
+}
+
+// Contains returns whether this name is in the set. Works for nil
+// sets too.
+func (ts TypeNameSet) Contains(val TypeName) bool {
+	if ts == nil {
+		return false
+	}
+	_, found := ts[val]
+	return found
+}
+
+// SetUnion returns a new set with all of the names in s1 or s2.
+func SetUnion(s1, s2 TypeNameSet) TypeNameSet {
+	var result TypeNameSet
+	for val := range s1 {
+		result = result.Add(val)
+	}
+	for val := range s2 {
+		result = result.Add(val)
+	}
+	return result
+}
+
+// StripUnusedDefinitions removes all types that aren't in roots or
+// referred to by the types in roots, for example types that are
+// generated as a byproduct of an allOf element.
+func StripUnusedDefinitions(
+	roots TypeNameSet,
+	definitions []TypeDefiner,
+) ([]TypeDefiner, error) {
+	// Collect all the reference sets for each type.
+	references := make(map[TypeName]TypeNameSet)
+
+	for _, def := range definitions {
+		references[*def.Name()] = def.Type().References()
+	}
+
+	graph := newReferenceGraph(roots, references)
+	connectedTypes := graph.connected()
+	var usedDefinitions []TypeDefiner
+	for _, def := range definitions {
+		if connectedTypes.Contains(*def.Name()) {
+			usedDefinitions = append(usedDefinitions, def)
+		}
+	}
+	return usedDefinitions, nil
+}
+
+// CollectResourceDefinitions returns a TypeNameSet of all of the
+// resource definitions in the definitions passed in.
+func CollectResourceDefinitions(definitions []TypeDefiner) TypeNameSet {
+	resources := make(TypeNameSet)
+	for _, def := range definitions {
+		if _, ok := def.(*ResourceDefinition); ok {
+			resources.Add(*def.Name())
+		}
+	}
+	return resources
+}
+
+func newReferenceGraph(roots TypeNameSet, references map[TypeName]TypeNameSet) *referenceGraph {
+	return &referenceGraph{
+		roots:      roots,
+		references: references,
+	}
+}
+
+type referenceGraph struct {
+	roots      TypeNameSet
+	references map[TypeName]TypeNameSet
+}
+
+// connected returns the set of types that are reachable from the
+// roots.
+func (c referenceGraph) connected() TypeNameSet {
+	// Make a non-nil set so we don't need to worry about passing it back down.
+	connectedTypes := make(TypeNameSet)
+	for node := range c.roots {
+		c.collectTypes(node, connectedTypes)
+	}
+	return connectedTypes
+}
+
+func (c referenceGraph) collectTypes(node TypeName, collected TypeNameSet) {
+	if collected.Contains(node) {
+		// We can stop here - we've already visited this node.
+		return
+	}
+	collected.Add(node)
+	for child := range c.references[node] {
+		c.collectTypes(child, collected)
+	}
+}

hack/generator/pkg/astmodel/strip_unused_test.go

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+package astmodel
+
+import (
+	"testing"
+
+	. "github.com/onsi/gomega"
+)
+
+const packagePath = "test.package/v1"
+
+func TestConnectionChecker_Avoids_Cycles(t *testing.T) {
+	g := NewGomegaWithT(t)
+	makeName := func(name string) TypeName {
+		return *NewTypeName(*NewPackageReference(packagePath), name)
+	}
+
+	makeSet := func(names ...string) TypeNameSet {
+		var typeNames []TypeName
+		for _, n := range names {
+			typeNames = append(typeNames, makeName(n))
+		}
+		return NewTypeNameSet(typeNames...)
+	}
+
+	roots := makeSet("res1", "res2")
+	references := map[TypeName]TypeNameSet{
+		makeName("G1"):   makeSet("G2"),
+		makeName("G2"):   makeSet("A"),
+		makeName("res1"): makeSet("A"),
+		makeName("A"):    makeSet("B", "C"),
+		makeName("B"):    nil,
+		makeName("C"):    makeSet("D"),
+		makeName("D"):    makeSet("A"), // cyclic
+	}
+
+	graph := newReferenceGraph(roots, references)
+	connectedSet := graph.connected()
+
+	g.Expect(connectedSet).To(Equal(makeSet("res1", "res2", "A", "B", "C", "D")))
+}

hack/generator/pkg/astmodel/struct_type.go

@@ -80,18 +80,17 @@ func (structType *StructType) RequiredImports() []*PackageReference {
 	return result
 }
 
-// References this type has to the given type
-func (structType *StructType) References(d *TypeName) bool {
-
+// References returns the set of all the types the fields referred to
+// by any field.
+func (structType *StructType) References() TypeNameSet {
+	var results TypeNameSet
 	for _, field := range structType.fields {
-		if field.FieldType().References(d) {
-			return true
+		for ref := range field.FieldType().References() {
+			results = results.Add(ref)
 		}
 	}
-
-	// For now, not considering functions in references on purpose
-
-	return false
+	// Not collecting types from functions deliberately.
+	return results
 }
 
 // Equals returns true if the passed type is a struct type with the same fields, false otherwise

hack/generator/pkg/astmodel/type.go

@@ -14,9 +14,10 @@ type Type interface {
 	// RequiredImports returns a list of packages required by this type
 	RequiredImports() []*PackageReference
 
-	// References determines if this type has a direct reference to the given definition name
-	// For example an Array of Persons references a Person
-	References(d *TypeName) bool
+	// References returns the names of all types that this type
+	// references. For example, an Array of Persons references a
+	// Person.
+	References() TypeNameSet
 
 	// AsType renders as a Go abstract syntax tree for a type
 	// (yes this says ast.Expr but that is what the Go 'ast' package uses for types)

hack/generator/pkg/astmodel/type_name.go

@@ -53,9 +53,12 @@ func (typeName *TypeName) AsType(codeGenerationContext *CodeGenerationContext) a
 	return ast.NewIdent(typeName.name)
 }
 
-// References indicates whether this Type includes any direct references to the given Type
-func (typeName *TypeName) References(d *TypeName) bool {
-	return typeName.Equals(d)
+// References returns a set containing this type name.
+func (typeName *TypeName) References() TypeNameSet {
+	if typeName == nil {
+		return nil
+	}
+	return NewTypeNameSet(*typeName)
 }
 
 // RequiredImports returns all the imports required for this definition

hack/generator/pkg/codegen/code_generator.go

@@ -81,6 +81,12 @@ func (generator *CodeGenerator) Generate(ctx context.Context) error {
 		return errors.Wrapf(err, "failed to filter generated definitions")
 	}
 
+	roots := astmodel.CollectResourceDefinitions(defs)
+	defs, err = astmodel.StripUnusedDefinitions(roots, defs)
+	if err != nil {
+		return errors.Wrapf(err, "failed to strip unused definitions")
+	}
+
 	packages, err := generator.CreatePackagesForDefinitions(defs)
 	if err != nil {
 		return errors.Wrapf(err, "failed to assign generated definitions to packages")

hack/generator/pkg/jsonast/jsonast_test.go

@@ -45,6 +45,18 @@ func runGoldenTest(t *testing.T, path string) {
 		t.Fatalf("could not produce nodes from scanner: %v", err)
 	}
 
+	// The golden files always generate a top-level Test type - mark
+	// that as the root.
+	roots := astmodel.NewTypeNameSet(*astmodel.NewTypeName(
+		*astmodel.NewPackageReference(
+			"github.com/Azure/k8s-infra/hack/generator/apis/test/v20200101"),
+		"Test",
+	))
+	defs, err = astmodel.StripUnusedDefinitions(roots, defs)
+	if err != nil {
+		t.Fatalf("could not strip unused types: %v", err)
+	}
+
 	// put all definitions in one file, regardless
 	// the package reference isn't really used here
 	fileDef := astmodel.NewFileDefinition(&defs[0].Name().PackageReference, defs...)

hack/generator/pkg/jsonast/testdata/AllOf/AllOf_generates_wrapper_type.golden

@@ -3,25 +3,6 @@
 // Code generated by k8s-infra-gen. DO NOT EDIT.
 package v20200101
 
-//Generated from: https://test.test/schemas/2020-01-01/test.json#/definitions/Bar
-type Bar struct {
-
-	// +kubebuilder:validation:Required
-	Size int `json:"size"`
-}
-
-//Generated from: https://test.test/schemas/2020-01-01/test.json#/definitions/Baz
-type Baz struct {
-
-	// +kubebuilder:validation:Required
-	Enabled bool `json:"enabled"`
-}
-
-//Generated from: https://test.test/schemas/2020-01-01/test.json#/definitions/Foo
-type Foo struct {
-	Name *string `json:"name"`
-}
-
 //Generated from: https://test.test/schemas/2020-01-01/test.json
 type Test struct {