Convert Direction from enum into a discriminated union

hack/generator/pkg/conversions/direction.go

@@ -5,12 +5,76 @@
 
 package conversions
 
+import (
+	"github.com/Azure/azure-service-operator/hack/generator/pkg/astmodel"
+)
+
 // Direction specifies the direction of conversion we're implementing with this function
-type Direction int
+type Direction interface {
+	// SelectString returns one of the provided strings, depending on the direction of conversion
+	SelectString(from string, to string) string
+	// SelectType returns one of the provided types, depending on the direction of conversion
+	SelectType(from astmodel.Type, to astmodel.Type) astmodel.Type
+	// WhenFrom will run the specified function only if the direction is "From", returning the current direction for chaining
+	WhenFrom(fn func()) Direction
+	// WhenTo will run the specified function only if the direction is "To", returning the current direction for chaining
+	WhenTo(fn func()) Direction
+}
 
-const (
+var (
 	// ConvertFrom indicates the conversion is from the passed 'other', populating the receiver with properties from the other
-	ConvertFrom = Direction(1)
+	ConvertFrom Direction = &ConvertFromDirection{}
 	// ConvertTo indicates the conversion is to the passed 'other', populating the other with properties from the receiver
-	ConvertTo = Direction(2)
+	ConvertTo Direction = &ConvertToDirection{}
 )
+
+type ConvertFromDirection struct{}
+
+var _ Direction = &ConvertFromDirection{}
+
+// SelectString returns the string for conversion FROM
+func (dir *ConvertFromDirection) SelectString(fromString string, _ string) string {
+	return fromString
+}
+
+// SelectType returns the type for conversion FROM
+func (dir *ConvertFromDirection) SelectType(fromType astmodel.Type, _ astmodel.Type) astmodel.Type {
+	return fromType
+}
+
+// WhenFrom will run the supplied function, returning this FROM direction for chaining
+func (dir *ConvertFromDirection) WhenFrom(fn func()) Direction {
+	fn()
+	return dir
+}
+
+// WhenTo will skip the supplied function, returning this FROM direction for chaining
+func (dir *ConvertFromDirection) WhenTo(_ func()) Direction {
+	// Nothing
+	return dir
+}
+
+type ConvertToDirection struct{}
+
+var _ Direction = &ConvertToDirection{}
+
+// SelectString returns the string for conversion TO
+func (dir *ConvertToDirection) SelectString(_ string, toValue string) string {
+	return toValue
+}
+
+// SelectType returns the type for conversion TO
+func (dir *ConvertToDirection) SelectType(_ astmodel.Type, toType astmodel.Type) astmodel.Type {
+	return toType
+}
+
+// WhenFrom will skip the supplied function, returning this TO direction for chaining
+func (dir *ConvertToDirection) WhenFrom(_ func()) Direction {
+	return dir
+}
+
+// WhenTo will run the supplied function, returning this TO direction for chaining
+func (dir *ConvertToDirection) WhenTo(fn func()) Direction {
+	fn()
+	return dir
+}

hack/generator/pkg/conversions/property_assignment_function.go

@@ -154,15 +154,10 @@ func (fn *PropertyAssignmentFunction) Equals(f astmodel.Function) bool {
 
 // AsFunc renders this function as an AST for serialization to a Go source file
 func (fn *PropertyAssignmentFunction) AsFunc(generationContext *astmodel.CodeGenerationContext, receiver astmodel.TypeName) *dst.FuncDecl {
-	var description string
-	switch fn.direction {
-	case ConvertFrom:
-		description = fmt.Sprintf("populates our %s from the provided source %s", receiver.Name(), fn.otherDefinition.Name().Name())
-	case ConvertTo:
-		description = fmt.Sprintf("populates the provided destination %s from our %s", fn.otherDefinition.Name().Name(), receiver.Name())
-	default:
-		panic(fmt.Sprintf("unexpected conversion direction %q", fn.direction))
-	}
+
+	description := fn.direction.SelectString(
+		fmt.Sprintf("populates our %s from the provided source %s", receiver.Name(), fn.otherDefinition.Name().Name()),
+		fmt.Sprintf("populates the provided destination %s from our %s", fn.otherDefinition.Name().Name(), receiver.Name()))
 
 	// We always use a pointer receiver so we can modify it
 	receiverType := astmodel.NewOptionalType(receiver).AsType(generationContext)
@@ -268,19 +263,14 @@ func (fn *PropertyAssignmentFunction) generateAssignments(
 // createConversions iterates through the properties on our receiver type, matching them up with
 // our other type and generating conversions where possible
 func (fn *PropertyAssignmentFunction) createConversions(receiver astmodel.TypeDefinition) error {
-	var sourceEndpoints map[string]ReadableConversionEndpoint
-	var destinationEndpoints map[string]WritableConversionEndpoint
+	// When converting FROM, otherDefinition.Type() is our source
+	// When converting TO, receiver.Type() is our source
+	// and conversely for our destination
+	sourceType := fn.direction.SelectType(fn.otherDefinition.Type(), receiver.Type())
+	destinationType := fn.direction.SelectType(receiver.Type(), fn.otherDefinition.Type())
 
-	switch fn.direction {
-	case ConvertFrom:
-		sourceEndpoints = fn.createReadableEndpoints(fn.otherDefinition.Type())
-		destinationEndpoints = fn.createWritableEndpoints(receiver.Type())
-	case ConvertTo:
-		sourceEndpoints = fn.createReadableEndpoints(receiver.Type())
-		destinationEndpoints = fn.createWritableEndpoints(fn.otherDefinition.Type())
-	default:
-		panic(fmt.Sprintf("unexpected conversion direction %q", fn.direction))
-	}
+	sourceEndpoints := fn.createReadableEndpoints(sourceType)
+	destinationEndpoints := fn.createWritableEndpoints(destinationType)
 
 	// Flag receiver and parameter names as used
 	fn.knownLocals.Add(fn.receiverName)
@@ -301,8 +291,8 @@ func (fn *PropertyAssignmentFunction) createConversions(receiver astmodel.TypeDe
 			return errors.Wrapf(
 				err,
 				"creating conversion to %s by %s",
-				sourceEndpoint,
-				destinationEndpoint)
+				destinationEndpoint,
+				sourceEndpoint)
 		} else if conv != nil {
 			// A conversion was created, keep it for later
 			fn.conversions[destinationName] = conv
@@ -377,13 +367,7 @@ func (fn *PropertyAssignmentFunction) createConversion(
 
 func nameOfPropertyAssignmentFunction(name astmodel.TypeName, direction Direction, idFactory astmodel.IdentifierFactory) string {
 	nameOfOtherType := idFactory.CreateIdentifier(name.Name(), astmodel.Exported)
-	switch direction {
-	case ConvertTo:
-		return "AssignPropertiesTo" + nameOfOtherType
-
-	case ConvertFrom:
-		return "AssignPropertiesFrom" + nameOfOtherType
-	}
-
-	panic(fmt.Sprintf("unexpected conversion direction %q", direction))
+	return direction.SelectString(
+		"AssignPropertiesFrom"+nameOfOtherType,
+		"AssignPropertiesTo"+nameOfOtherType)
 }