From 7c2de982ec4aaba0a165d5de6e62ed07d29d3879 Mon Sep 17 00:00:00 2001
From: Bevan Arps <bevan.arps@microsoft.com>
Date: Thu, 10 Jun 2021 16:08:48 +1200
Subject: [PATCH] Simplify conversion context

---
 .../pkg/codegen/storage/storage_type_factory.go        |  7 +++----
 .../pkg/conversions/property_assignment_function.go    |  8 ++++++--
 .../pkg/conversions/property_conversion_context.go     | 10 ++++++++--
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/hack/generator/pkg/codegen/storage/storage_type_factory.go b/hack/generator/pkg/codegen/storage/storage_type_factory.go
index e355788c407..b7be238993d 100644
--- 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
index 38586121992..281f3315fef 100644
--- a/hack/generator/pkg/conversions/property_assignment_function.go
+++ b/hack/generator/pkg/conversions/property_assignment_function.go
@@ -60,7 +60,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)
 	if err != nil {
@@ -85,7 +87,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)
 	if err != nil {
diff --git a/hack/generator/pkg/conversions/property_conversion_context.go b/hack/generator/pkg/conversions/property_conversion_context.go
index cb86ea90ad2..268b0f2f2f9 100644
--- 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 {