From eb239f423e39b464d425bd9750d9c6238409fd5f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergen=20Yal=C3=A7=C4=B1n?= <yalcinsergen97@gmail.com>
Date: Thu, 21 Dec 2023 15:44:56 +0300
Subject: [PATCH] Add reference fields to the InitProvider
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Sergen Yalçın <yalcinsergen97@gmail.com>
---
 pkg/types/builder.go | 21 +++++++++++++++------
 pkg/types/field.go   |  6 +++---
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/pkg/types/builder.go b/pkg/types/builder.go
index 3c476fe8..6b9aa4a8 100644
--- a/pkg/types/builder.go
+++ b/pkg/types/builder.go
@@ -274,7 +274,7 @@ func (g *Builder) buildSchema(f *Field, cfg *config.Resource, names []string, r
 				if paramType.Underlying().String() != emptyStruct {
 					field := types.NewField(token.NoPos, g.Package, f.Name.Camel, types.NewSlice(paramType), false)
 					r.addParameterField(f, field)
-					r.addInitField(f, field, g.Package)
+					r.addInitField(f, field, g, nil)
 				}
 			default:
 				if paramType == nil {
@@ -400,7 +400,7 @@ func (r *resource) addParameterField(f *Field, field *types.Var) {
 	r.paramFields = append(r.paramFields, field)
 }
 
-func (r *resource) addInitField(f *Field, field *types.Var, pkg *types.Package) {
+func (r *resource) addInitField(f *Field, field *types.Var, g *Builder, typeNames *types.TypeName) {
 	// If the field is not an init field, we don't add it.
 	if !f.isInit() {
 		return
@@ -410,10 +410,14 @@ func (r *resource) addInitField(f *Field, field *types.Var, pkg *types.Package)
 
 	// If the field is a nested type, we need to add it as the init type.
 	if f.InitType != nil {
-		field = types.NewField(token.NoPos, pkg, f.Name.Camel, f.InitType, false)
+		field = types.NewField(token.NoPos, g.Package, f.Name.Camel, f.InitType, false)
 	}
 
 	r.initFields = append(r.initFields, field)
+
+	if f.Reference != nil {
+		r.addReferenceFields(g, typeNames, f, true)
+	}
 }
 
 func (r *resource) addObservationField(f *Field, field *types.Var) {
@@ -429,10 +433,15 @@ func (r *resource) addObservationField(f *Field, field *types.Var) {
 	r.obsTags = append(r.obsTags, fmt.Sprintf(`json:"%s" tf:"%s"`, f.JSONTag, f.TFTag))
 }
 
-func (r *resource) addReferenceFields(g *Builder, paramName *types.TypeName, field *Field) {
+func (r *resource) addReferenceFields(g *Builder, paramName *types.TypeName, field *Field, isInit bool) {
 	refFields, refTags := g.generateReferenceFields(paramName, field)
-	r.paramTags = append(r.paramTags, refTags...)
-	r.paramFields = append(r.paramFields, refFields...)
+	if isInit {
+		r.initTags = append(r.initTags, refTags...)
+		r.initFields = append(r.initFields, refFields...)
+	} else {
+		r.paramTags = append(r.paramTags, refTags...)
+		r.paramFields = append(r.paramFields, refFields...)
+	}
 }
 
 // generateTypeName generates a unique name for the type if its original name
diff --git a/pkg/types/field.go b/pkg/types/field.go
index 89e2df08..3ea21d86 100644
--- a/pkg/types/field.go
+++ b/pkg/types/field.go
@@ -351,11 +351,11 @@ func (f *Field) AddToResource(g *Builder, r *resource, typeNames *TypeNames, add
 			f.TFTag = strings.TrimSuffix(f.TFTag, ",omitempty")
 		}
 		r.addParameterField(f, field)
-		r.addInitField(f, field, g.Package)
+		r.addInitField(f, field, g, typeNames.InitTypeName)
 	}
 
 	if f.Reference != nil {
-		r.addReferenceFields(g, typeNames.ParameterTypeName, f)
+		r.addReferenceFields(g, typeNames.ParameterTypeName, f, false)
 	}
 
 	// Note(lsviben): All fields are optional because observation fields are
@@ -401,7 +401,7 @@ func (f *Field) AddToResource(g *Builder, r *resource, typeNames *TypeNames, add
 // an earlier step, so they cannot be included as well. Plus probably they
 // should also not change for Create and Update steps.
 func (f *Field) isInit() bool {
-	return !f.Identifier && f.Reference == nil && (f.TFTag != "-" || f.Injected)
+	return !f.Identifier && f.TFTag != "-"
 }
 
 func getDescription(s string) string {