Fix ordering bug for OneOf processing (#4320)

* Removed unused field

* Allow AsSlice() to be used from a ReadonlyTypeSet

* Code gardening

* Ensure OneOf merging is not order sensitive

* Add test for new case

* Add new test helper

* Code formatting

v2/tools/generator/internal/astmodel/oneof_type.go

@@ -18,7 +18,6 @@ import (
 type OneOfType struct {
 	swaggerName           string        // Name of the OneOf as defined in the Swagger file
 	propertyObjects       []*ObjectType // Object definitions used to specify the properties held by this OneOf. May be empty.
-	options               TypeNameSet   // References to the type names of the options for this OneOf. May be empty.
 	types                 TypeSet       // Set of all possible types
 	discriminatorProperty string        // Identifies the discriminatorProperty property
 	discriminatorValue    string        // Discriminator value used to identify this subtype
@@ -31,7 +30,6 @@ func NewOneOfType(name string, types ...Type) *OneOfType {
 	return &OneOfType{
 		swaggerName: name,
 		types:       MakeTypeSet(types...),
-		options:     NewTypeNameSet(),
 	}
 }
 

v2/tools/generator/internal/astmodel/type_set.go

@@ -17,6 +17,7 @@ type ReadonlyTypeSet interface {
 	ForEachError(func(t Type, ix int) error) error
 	Len() int
 	Single() (Type, bool)
+	AsSlice() []Type
 }
 
 var _ ReadonlyTypeSet = TypeSet{}

v2/tools/generator/internal/codegen/pipeline/assemble_oneof.go

@@ -108,9 +108,12 @@ func (oa *oneOfAssembler) assemble(name astmodel.InternalTypeName) error {
 
 // assemblePair updates a pair of nodes that have a relationship and recursively invokes itself to traverse the entire
 // tree of referenced imports.
-// ancestor is the parent or super-type. It may be a less specialised OneOf, a root OneOf, an Object or an AllOf
+// ancestor is the parent or super-type. It may be a less specialised OneOf, a root OneOf, an Object or an AllOf.
 // descendent is the child or subtype. It will be an intermediate or leaf OneOf.
-func (oa *oneOfAssembler) assemblePair(ancestor astmodel.InternalTypeName, descendant astmodel.InternalTypeName) error {
+func (oa *oneOfAssembler) assemblePair(
+	ancestor astmodel.InternalTypeName,
+	descendant astmodel.InternalTypeName,
+) error {
 	// Remove any direct reference to the parent from the leaf
 	err := oa.removeParentReferenceFromLeaf(ancestor, descendant)
 	if err != nil {

v2/tools/generator/internal/codegen/pipeline/convert_allof_and_oneof_to_objects.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/pkg/errors"
 	"golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
 
 	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/astmodel"
 )
@@ -884,35 +885,35 @@ func (synthesizer) handleMapObject(leftMap *astmodel.MapType, rightObj *astmodel
 // makes an ObjectType for an AllOf type
 // We're all about synthesizing a new type, so we resolve TypeNames here
 func (s synthesizer) allOfObject(allOf *astmodel.AllOfType) (astmodel.Type, error) {
-	types := allOf.Types()
-	toMerge := make([]astmodel.Type, types.Len())
-	types.ForEach(func(t astmodel.Type, i int) {
-		toMerge[i] = t
-	})
-
+	toMerge := allOf.Types().AsSlice()
 	return s.allOfSlice(toMerge)
 }
 
 func (s synthesizer) allOfSlice(types []astmodel.Type) (astmodel.Type, error) {
-	toMerge := make([]astmodel.Type, len(types))
-	foundName := false
-	for i, t := range types {
-		// if we find a type name, resolve it to the underlying type
-		if tn, ok := astmodel.AsInternalTypeName(t); ok {
-			if def, ok := s.defs[tn]; ok {
-				toMerge[i] = def.Type()
-				foundName = true
-				continue
+	toMerge := s.simplifyAllOfTypeNames(types)
+
+	// When an ObjectType is merged with a OneOfType, two things happen. The properties on the object type are pushed
+	// down to all the "leaves" of the OneOfType, and the OneOfType is converted to an ObjectType.
+	// If we later try to merge another ObjectType with the OneOf-ObjectType, the properties from the new object end up
+	// merged onto the OneOf-ObjectType itself instead of being pushed down to the leaves.
+	// To avoid this from happening, we need to merge all the ObjectTypes first, before we merge them with the OneOf.
+	// We achieve this by ordering the types so that all ObjectTypes come first.
+	slices.SortFunc(
+		toMerge,
+		func(left astmodel.Type, right astmodel.Type) int {
+			_, leftIsObject := left.(*astmodel.ObjectType)
+			_, rightIsObject := right.(*astmodel.ObjectType)
+
+			if leftIsObject && !rightIsObject {
+				return -1
 			}
-		}
 
-		toMerge[i] = t
-	}
+			if !leftIsObject && rightIsObject {
+				return 1
+			}
 
-	if foundName {
-		// Found a type name, recursive call in case there are more
-		return s.allOfSlice(toMerge)
-	}
+			return 0
+		})
 
 	var result astmodel.Type = astmodel.AnyType
 	for _, t := range toMerge {
@@ -926,6 +927,32 @@ func (s synthesizer) allOfSlice(types []astmodel.Type) (astmodel.Type, error) {
 	return result, nil
 }
 
+// simplifyAllOfTypeNames converts any type names in the allOf slice to their underlying types
+// so that we can merge everything.
+// AllOf types will reference another type to "import" all of its properties.
+// It's possible for a TypeDefinition to be an alias for another, so we recurse to follow any chains.
+func (s synthesizer) simplifyAllOfTypeNames(types []astmodel.Type) []astmodel.Type {
+	foundName := false
+	result := make([]astmodel.Type, len(types))
+	for i, t := range types {
+		if tn, ok := astmodel.AsInternalTypeName(t); ok {
+			if def, ok := s.defs[tn]; ok {
+				result[i] = def.Type()
+				foundName = true
+				continue
+			}
+		}
+
+		result[i] = t
+	}
+
+	if foundName {
+		return s.simplifyAllOfTypeNames(result)
+	}
+
+	return result
+}
+
 func (s synthesizer) handleFlaggedType(left *astmodel.FlaggedType, right astmodel.Type) (astmodel.Type, error) {
 	// Intersect the content and retain the flags
 	internal, err := s.intersectTypes(left.Element(), right)

v2/tools/generator/internal/codegen/pipeline/convert_allof_and_oneof_to_objects_test.go

@@ -9,11 +9,11 @@ import (
 	"context"
 	"testing"
 
-	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/test"
+	. "github.com/onsi/gomega"
 
+	"github.com/Azure/azure-service-operator/v2/internal/util/to"
 	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/astmodel"
-
-	. "github.com/onsi/gomega"
+	"github.com/Azure/azure-service-operator/v2/tools/generator/internal/test"
 )
 
 func makeSynth(definitions ...astmodel.TypeDefinition) synthesizer {
@@ -713,3 +713,105 @@ func TestConversionOfSequentialOneOf_ReturnsExpectedResults(t *testing.T) {
 		test.AssertDefinitionHasExpectedShape(t, def.Name().Name(), def)
 	}
 }
+
+// Checking a scenario discovered while importing the Kusto group.
+// When a resource spec contains an allOf that in turn contains a oneOf, we need all the properties pushed down
+// to the oneof leaves
+func TestConversionOfAllOf_WhenContainingOneOf_ReturnsExpectedResult(t *testing.T) {
+	t.Parallel()
+
+	g := NewGomegaWithT(t)
+
+	idFactory := astmodel.NewIdentifierFactory()
+
+	readwriteDatabase := astmodel.MakeTypeDefinition(
+		astmodel.MakeInternalTypeName(test.Pkg2020, "ReadWriteDatabase"),
+		astmodel.NewObjectType().
+			WithProperties(
+				astmodel.NewPropertyDefinition(
+					"KeyVaultURL",
+					"keyVaultUrl",
+					astmodel.StringType),
+				astmodel.NewPropertyDefinition(
+					"HotCachePeriod",
+					"hotCachePeriod",
+					astmodel.StringType),
+				astmodel.NewPropertyDefinition(
+					"Kind",
+					"kind",
+					astmodel.NewEnumType(astmodel.StringType, astmodel.MakeEnumValue("ReadWriteDatabase", "ReadWriteDatabase"))),
+				astmodel.NewPropertyDefinition(
+					"Location",
+					"location",
+					astmodel.StringType)),
+	)
+
+	readonlyFollowingDatabase := astmodel.MakeTypeDefinition(
+		astmodel.MakeInternalTypeName(test.Pkg2020, "ReadOnlyFollowingDatabase"),
+		astmodel.NewObjectType().
+			WithProperties(
+				astmodel.NewPropertyDefinition(
+					"DatabaseShareOrigin",
+					"databaseShareOrigin",
+					astmodel.StringType),
+				astmodel.NewPropertyDefinition(
+					"HotCachePeriod",
+					"hotCachePeriod",
+					astmodel.StringType),
+				astmodel.NewPropertyDefinition(
+					"Kind",
+					"kind",
+					astmodel.NewEnumType(astmodel.StringType, astmodel.MakeEnumValue("ReadOnlyFollowingDatabase", "ReadOnlyFollowingDatabase"))),
+				astmodel.NewPropertyDefinition(
+					"Location",
+					"location",
+					astmodel.StringType)),
+	)
+
+	database := astmodel.MakeTypeDefinition(
+		astmodel.MakeInternalTypeName(test.Pkg2020, "Database"),
+		astmodel.NewOneOfType(
+			"database",
+			readwriteDatabase.Name(),
+			readonlyFollowingDatabase.Name()).
+			WithDiscriminatorProperty("Kind"),
+	)
+
+	clusters_database := astmodel.MakeTypeDefinition(
+		astmodel.MakeInternalTypeName(test.Pkg2020, "Clusters_Database"),
+		astmodel.NewResourceType(
+			astmodel.NewAllOfType(
+				database.Name(),
+				astmodel.NewObjectType().
+					WithProperty(
+						astmodel.NewPropertyDefinition(
+							"Name",
+							"name",
+							astmodel.NewValidatedType(
+								astmodel.StringType,
+								astmodel.StringValidations{
+									MaxLength: to.Ptr(int64(96)),
+								}))),
+				astmodel.NewObjectType().
+					WithProperty(
+						astmodel.NewPropertyDefinition(
+							"Name",
+							"name",
+							astmodel.StringType))),
+			astmodel.NewObjectType(),
+		))
+
+	defs := astmodel.MakeTypeDefinitionSetFromDefinitions(
+		readwriteDatabase,
+		readonlyFollowingDatabase,
+		database,
+		clusters_database)
+
+	state := NewState(defs)
+	stage := ConvertAllOfAndOneOfToObjects(idFactory)
+
+	finalState, err := stage.Run(context.TODO(), state)
+	g.Expect(err).To(BeNil())
+
+	test.AssertDefinitionsHaveExpectedShapes(t, "structure.txt", finalState.Definitions())
+}

v2/tools/generator/internal/codegen/pipeline/testdata/TestConversionOfAllOf_WhenContainingOneOf_ReturnsExpectedResult/structure.txt.golden

@@ -0,0 +1,28 @@
+github.com/Azure/azure-service-operator/testing/person/v20200101
+----------------------------------------------------------------
+Clusters_Database: Resource
+├── Spec: Object (2 properties)
+│   ├── Flag 0: oneof
+│   ├── ReadOnlyFollowing: *ReadOnlyFollowingDatabase
+│   └── ReadWrite: *ReadWriteDatabase
+└── Status: Object (0 properties)
+Database: Object (2 properties)
+├── Flag 0: oneof
+├── ReadOnlyFollowing: *ReadOnlyFollowingDatabase
+└── ReadWrite: *ReadWriteDatabase
+ReadOnlyFollowingDatabase: Object (5 properties)
+├── DatabaseShareOrigin: string
+├── HotCachePeriod: string
+├── Kind: Enum (1 value)
+│   └── ReadOnlyFollowingDatabase
+├── Location: string
+└── Name: Validated<string> (1 rule)
+    └── Rule 0: MaxLength: 96
+ReadWriteDatabase: Object (5 properties)
+├── HotCachePeriod: string
+├── KeyVaultURL: string
+├── Kind: Enum (1 value)
+│   └── ReadWriteDatabase
+├── Location: string
+└── Name: Validated<string> (1 rule)
+    └── Rule 0: MaxLength: 96

v2/tools/generator/internal/test/asserts.go

@@ -89,6 +89,20 @@ func AssertDefinitionHasExpectedShape(
 	t *testing.T,
 	filename string,
 	def astmodel.TypeDefinition,
+) {
+	defs := make(astmodel.TypeDefinitionSet)
+	defs.Add(def)
+	AssertDefinitionsHaveExpectedShapes(t, filename, defs)
+}
+
+// AssertDefinitionsHaveExpectedShapes fails the test if the given definition does not have the expected shape.
+// t is the current test.
+// filename is the name of the golden file to write.
+// def is the definition to be asserted.
+func AssertDefinitionsHaveExpectedShapes(
+	t *testing.T,
+	filename string,
+	defs astmodel.TypeDefinitionSet,
 ) {
 	t.Helper()
 	g := goldie.New(t)
@@ -98,9 +112,6 @@ func AssertDefinitionHasExpectedShape(
 		t.Fatalf("Unable to configure goldie output folder %s", err)
 	}
 
-	defs := make(astmodel.TypeDefinitionSet)
-	defs.Add(def)
-
 	buf := &bytes.Buffer{}
 	report := reporting.NewTypeCatalogReport(defs)
 	err = report.WriteTo(buf)