Refactor step3codegen package (mainly TypeGenerator)

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step2intermediate/MoreTypes.kt

@@ -1,48 +1,46 @@
 package com.virtuslab.pulumikotlin.codegen.step2intermediate
 
-import com.squareup.kotlinpoet.ClassName
+import com.pulumi.core.Output
+import com.pulumi.kotlin.ConvertibleToJava
+import com.pulumi.kotlin.PulumiTagMarker
+import com.pulumi.kotlin.options.CustomResourceOptions
+import com.pulumi.kotlin.options.CustomResourceOptionsBuilder
 import com.squareup.kotlinpoet.MemberName
 import com.squareup.kotlinpoet.MemberName.Companion.member
-import com.squareup.kotlinpoet.ParameterizedTypeName
-import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
-import com.squareup.kotlinpoet.TypeName
+import com.squareup.kotlinpoet.asClassName
+import com.squareup.kotlinpoet.asTypeName
 
 object MoreTypes {
 
     object Kotlin {
         object Pulumi {
-            fun toJava(): MemberName {
-                return MemberName("com.pulumi.kotlin", "toJava")
-            }
+            fun toJavaExtensionMethod() = MemberName("com.pulumi.kotlin", "toJava")
 
-            fun applySuspend(): MemberName {
-                return MemberName("com.pulumi.kotlin", "applySuspend")
-            }
+            fun toKotlinExtensionMethod() = MemberName("com.pulumi.kotlin", "toKotlin")
 
-            fun ConvertibleToJava(type: TypeName): ParameterizedTypeName =
-                ClassName("com.pulumi.kotlin", "ConvertibleToJava").parameterizedBy(type)
-        }
+            fun applySuspendExtensionMethod() = MemberName("com.pulumi.kotlin", "applySuspend")
+
+            fun applyValueExtensionMethod() = MemberName("com.pulumi.kotlin", "applyValue")
+
+            fun pulumiDslMarkerAnnotation() = PulumiTagMarker::class.asClassName()
+
+            fun convertibleToJavaClass() = ConvertibleToJava::class.asClassName()
+
+            fun customResourceOptionsClass() = CustomResourceOptions::class.asClassName()
 
-        fun Pair(leftType: TypeName, rightType: TypeName): ParameterizedTypeName {
-            return ClassName("kotlin", "Pair").parameterizedBy(leftType, rightType)
+            fun customResourceOptionsBuilderClass() = CustomResourceOptionsBuilder::class.asClassName()
         }
+
+        fun coroutinesFutureAwaitExtensionMethod() = MemberName("kotlinx.coroutines.future", "await")
+
+        fun pairClass() = Pair::class.asTypeName()
     }
 
     object Java {
         object Pulumi {
+            fun outputOfMethod() = outputClass().member("of")
 
-            object Output {
-                val of: MemberName
-                    get() = Output().member("of")
-            }
-
-            fun Output(): ClassName {
-                return ClassName("com.pulumi.core", "Output")
-            }
-
-            fun Output(type: TypeName): ParameterizedTypeName {
-                return ClassName("com.pulumi.core", "Output").parameterizedBy(type)
-            }
+            fun outputClass() = Output::class.asClassName()
         }
     }
 }

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step3codegen/FieldTypes.kt

@@ -2,11 +2,11 @@ package com.virtuslab.pulumikotlin.codegen.step3codegen
 
 import com.squareup.kotlinpoet.ParameterizedTypeName
 import com.squareup.kotlinpoet.TypeName
-import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Java.Pulumi.outputClass
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.ReferencedType
+import com.virtuslab.pulumikotlin.codegen.step3codegen.KotlinPoetExtensions.parameterizedBy
 
 sealed class FieldType<T : ReferencedType> {
-
     abstract val type: T
 
     abstract fun toTypeName(): TypeName
@@ -20,14 +20,20 @@ data class NormalField<T : ReferencedType>(override val type: T, val mappingCode
 
 data class OutputWrappedField<T : ReferencedType>(override val type: T) : FieldType<T>() {
     override fun toTypeName(): ParameterizedTypeName {
-        return MoreTypes.Java.Pulumi.Output(type.toTypeName())
+        return outputClass().parameterizedBy(type)
     }
 }
 
 data class Field<T : ReferencedType>(
     val name: String,
     val fieldType: FieldType<T>,
     val required: Boolean,
-    val overloads: List<FieldType<*>> = emptyList(),
+    val overloads: List<FieldType<ReferencedType>> = emptyList(),
     val kDoc: KDoc,
-)
+) {
+    fun toTypeName(): TypeName =
+        fieldType.toTypeName().copy(nullable = !required)
+
+    fun toNullableTypeName(): TypeName =
+        fieldType.toTypeName().copy(nullable = true)
+}

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step3codegen/KotlinPoetExtensions.kt

@@ -0,0 +1,40 @@
+package com.virtuslab.pulumikotlin.codegen.step3codegen
+
+import com.squareup.kotlinpoet.ClassName
+import com.squareup.kotlinpoet.CodeBlock
+import com.squareup.kotlinpoet.FileSpec
+import com.squareup.kotlinpoet.FunSpec
+import com.squareup.kotlinpoet.MemberName
+import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
+import com.squareup.kotlinpoet.PropertySpec
+import com.squareup.kotlinpoet.TypeSpec
+import com.virtuslab.pulumikotlin.codegen.expressions.Code
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.Type
+
+object KotlinPoetExtensions {
+    fun TypeSpec.Builder.addFunctions(vararg funSpecs: FunSpec) = addFunctions(funSpecs.toList())
+
+    fun TypeSpec.Builder.addProperties(vararg propertySpecs: PropertySpec) = addProperties(propertySpecs.toList())
+
+    fun FileSpec.Builder.addImport(memberName: MemberName) = addImport(memberName.packageName, memberName.simpleName)
+
+    fun FileSpec.Builder.addImports(vararg memberNames: MemberName) =
+        apply {
+            memberNames.forEach {
+                addImport(it)
+            }
+        }
+
+    fun FileSpec.Builder.addTypes(vararg typeSpecs: TypeSpec) = addTypes(typeSpecs.toList())
+
+    fun FileSpec.Builder.addTypes(types: Iterable<TypeSpec>) =
+        apply {
+            types.forEach {
+                addType(it)
+            }
+        }
+
+    fun ClassName.parameterizedBy(vararg types: Type) = parameterizedBy(types.map { it.toTypeName() })
+
+    fun CodeBlock.Builder.add(code: Code) = add(code.toCodeBlock().toKotlinPoetCodeBlock())
+}

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step3codegen/KotlinPoetPatterns.kt

@@ -1,6 +1,9 @@
 package com.virtuslab.pulumikotlin.codegen.step3codegen
 
 import com.squareup.kotlinpoet.CodeBlock
+import com.squareup.kotlinpoet.FunSpec
+import com.squareup.kotlinpoet.KModifier
+import com.squareup.kotlinpoet.KModifier.SUSPEND
 import com.squareup.kotlinpoet.LIST
 import com.squareup.kotlinpoet.LambdaTypeName
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
@@ -10,26 +13,60 @@ import com.virtuslab.pulumikotlin.codegen.expressions.Assignment
 import com.virtuslab.pulumikotlin.codegen.expressions.CustomExpression
 import com.virtuslab.pulumikotlin.codegen.expressions.Expression
 import com.virtuslab.pulumikotlin.codegen.expressions.add
+import com.virtuslab.pulumikotlin.codegen.expressions.callLet
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.ComplexType
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.ReferencedComplexType
 
 typealias MappingCode = (Expression) -> Expression
 
 object KotlinPoetPatterns {
 
+    /**
+     * Example output:
+     *
+     * ```
+     * List<suspend InnerType.() -> Unit>
+     * ```
+     */
     fun listOfLambdas(innerType: TypeName): TypeName =
         LIST.parameterizedBy(builderLambda(innerType))
 
+    /**
+     * @see listOfLambdas
+     */
+    fun listOfLambdas(innerType: ComplexType): TypeName =
+        listOfLambdas(innerType.toBuilderTypeName())
+
+    /**
+     * @see listOfLambdas
+     */
     fun listOfLambdas(innerType: ReferencedComplexType): TypeName =
         listOfLambdas(innerType.toBuilderTypeName())
 
-    fun builderLambda(innerType: ReferencedComplexType): TypeName =
+    /**
+     * Example output:
+     *
+     * ```
+     * suspend InnerType.() -> Unit
+     * ```
+     */
+    fun builderLambda(innerType: ComplexType): TypeName =
         builderLambda(innerType.toBuilderTypeName())
 
+    /**
+     * @see builderLambda
+     */
     fun builderLambda(innerType: TypeName): TypeName =
         LambdaTypeName
             .get(receiver = innerType, returnType = UNIT)
             .copy(suspending = true)
 
+    /**
+     * @see builderLambda
+     */
+    fun builderLambda(innerType: ReferencedComplexType): TypeName =
+        builderLambda(innerType.toBuilderTypeName())
+
     data class BuilderSettingCodeBlock(val mappingCode: MappingCode? = null, val code: String, val args: List<Any?>) {
         fun withMappingCode(mappingCode: MappingCode): BuilderSettingCodeBlock {
             return copy(mappingCode = mappingCode)
@@ -54,4 +91,65 @@ object KotlinPoetPatterns {
             fun create(code: String, vararg args: Any?) = BuilderSettingCodeBlock(null, code, args.toList())
         }
     }
+
+    /**
+     * Example output:
+     *
+     * ```
+     * suspend fun name(vararg argument: ParameterType): Unit {
+     *    val toBeMapped = something.toList().somethingElse()
+     *    val mapped = mappingCode(toBeMapped)
+     *    this.field = mapped
+     * }
+     * ```
+     */
+    fun builderPattern(
+        name: String,
+        parameterType: TypeName,
+        kDoc: KDoc,
+        codeBlock: BuilderSettingCodeBlock,
+        parameterModifiers: List<KModifier> = emptyList(),
+    ): FunSpec {
+        return FunSpec
+            .builder(name)
+            .addModifiers(SUSPEND)
+            .addParameter(
+                "argument",
+                parameterType,
+                parameterModifiers,
+            )
+            .addCode(codeBlock.toCodeBlock(name))
+            .addDocsToBuilderMethod(kDoc, "argument")
+            .build()
+    }
+
+    /**
+     * Example output:
+     *
+     * ```
+     * val toBeMapped = something.toList().somethingElse()
+     * val mapped = mappingCode(toBeMapped)
+     * this.field = mapped
+     * ```
+     */
+    fun mappingCodeBlock(
+        field: NormalField<*>,
+        required: Boolean,
+        name: String,
+        code: String,
+        vararg args: Any?,
+    ): CodeBlock {
+        val expression = CustomExpression("toBeMapped").callLet(!required) { arg -> field.mappingCode(arg) }
+        return CodeBlock.builder()
+            .addStatement("val toBeMapped = $code", *args)
+            .add(Assignment("mapped", expression))
+            .addStatement("")
+            .addStatement("this.%N = mapped", name)
+            .build()
+    }
+
+    fun FunSpec.Builder.addDocsToBuilderMethod(kDoc: KDoc, paramName: String) = apply {
+        addDocs("@param $paramName ${kDoc.description.orEmpty()}")
+        addDeprecationWarningIfAvailable(kDoc)
+    }
 }

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step3codegen/functions/FunctionGenerator.kt

@@ -23,9 +23,11 @@ import com.virtuslab.pulumikotlin.codegen.step2intermediate.Direction.Output
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.FunctionType
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.LanguageType.Java
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.LanguageType.Kotlin
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Kotlin.coroutinesFutureAwaitExtensionMethod
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.NamingFlags
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.Subject.Function
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.Subject.Resource
+import com.virtuslab.pulumikotlin.codegen.step3codegen.KotlinPoetExtensions.addImport
 import com.virtuslab.pulumikotlin.codegen.step3codegen.KotlinPoetPatterns.builderLambda
 import com.virtuslab.pulumikotlin.codegen.step3codegen.addDeprecationWarningIfAvailable
 import com.virtuslab.pulumikotlin.codegen.step3codegen.addDocs
@@ -53,7 +55,7 @@ object FunctionGenerator {
                         name.toFunctionGroupObjectName(namingFlags),
                     )
                     .addType(objectSpecBuilder.build())
-                    .addImport("kotlinx.coroutines.future", "await")
+                    .addImport(coroutinesFutureAwaitExtensionMethod())
                     .build()
 
                 listOf(fileSpec)