Simplify conversion context

Azure · Jun 14, 2021 · 8a5d3d7 · 8a5d3d7
1 parent ec6ddbf
commit 8a5d3d7
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/hack/generator/pkg/codegen/storage/storage_type_factory.go b/hack/generator/pkg/codegen/storage/storage_type_factory.go
@@ -232,14 +232,13 @@ func (f *StorageTypeFactory) injectConversions(definition astmodel.TypeDefinitio
 	knownTypes.AddTypes(f.inputTypes.Except(f.outputTypes))
 	knownTypes.AddTypes(f.outputTypes)
 
-	conversionFromContext := conversions.NewStorageConversionContext(knownTypes, conversions.ConvertFrom, f.idFactory)
-	assignFromFn, err := conversions.NewPropertyAssignmentFromFunction(definition, nextDef, f.idFactory, conversionFromContext)
+	conversionContext := conversions.NewStorageConversionContext(knownTypes, f.idFactory)
+	assignFromFn, err := conversions.NewPropertyAssignmentFromFunction(definition, nextDef, f.idFactory, conversionContext)
 	if err != nil {
 		return nil, errors.Wrapf(err, "creating PropertyAssignmentFrom() function for %q", name)
 	}
 
-	conversionToContext := conversions.NewStorageConversionContext(knownTypes, conversions.ConvertTo, f.idFactory)
-	assignToFn, err := conversions.NewPropertyAssignmentToFunction(definition, nextDef, f.idFactory, conversionToContext)
+	assignToFn, err := conversions.NewPropertyAssignmentToFunction(definition, nextDef, f.idFactory, conversionContext)
 	if err != nil {
 		return nil, errors.Wrapf(err, "creating PropertyAssignmentTo() function for %q", name)
 	}

diff --git a/hack/generator/pkg/conversions/property_assignment_function.go b/hack/generator/pkg/conversions/property_assignment_function.go
@@ -62,7 +62,9 @@ func NewPropertyAssignmentFromFunction(
 		knownLocals:     astmodel.NewKnownLocalsSet(idFactory),
 	}
 
-	result.conversionContext = conversionContext.WithFunctionName(result.Name()).WithKnownLocals(result.knownLocals)
+	result.conversionContext = conversionContext.WithFunctionName(result.Name()).
+		WithKnownLocals(result.knownLocals).
+		WithDirection(ConvertFrom)
 
 	err := result.createConversions(receiver, conversionContext.Types())
 	if err != nil {
@@ -87,7 +89,9 @@ func NewPropertyAssignmentToFunction(
 		knownLocals:     astmodel.NewKnownLocalsSet(idFactory),
 	}
 
-	result.conversionContext = conversionContext.WithFunctionName(result.Name()).WithKnownLocals(result.knownLocals)
+	result.conversionContext = conversionContext.WithFunctionName(result.Name()).
+		WithKnownLocals(result.knownLocals).
+		WithDirection(ConvertTo)
 
 	err := result.createConversions(receiver, conversionContext.Types())
 	if err != nil {

diff --git a/hack/generator/pkg/conversions/property_conversion_context.go b/hack/generator/pkg/conversions/property_conversion_context.go
@@ -24,10 +24,9 @@ type PropertyConversionContext struct {
 }
 
 // NewStorageConversionContext creates a new instance of a PropertyConversionContext
-func NewStorageConversionContext(types astmodel.Types, direction Direction, idFactory astmodel.IdentifierFactory) *PropertyConversionContext {
+func NewStorageConversionContext(types astmodel.Types, idFactory astmodel.IdentifierFactory) *PropertyConversionContext {
 	return &PropertyConversionContext{
 		types:       types,
-		direction:   direction,
 		idFactory:   idFactory,
 		knownLocals: astmodel.NewKnownLocalsSet(idFactory),
 	}
@@ -57,6 +56,13 @@ func (c *PropertyConversionContext) WithKnownLocals(knownLocals *astmodel.KnownL
 	return result
 }
 
+// WithDirection returns a new context with the specified direction
+func (c *PropertyConversionContext) WithDirection( dir Direction) *PropertyConversionContext {
+	result := c.clone()
+	result.direction = dir
+	return result
+}
+
 // NestedContext returns a new context with a cloned knownLocals so that nested blocks
 // can declare and reuse locals independently of each other.
 func (c *PropertyConversionContext) NestedContext() *PropertyConversionContext {