Refactor step3codegen package (mainly TypeGenerator)

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/expressions/Expressions.kt

@@ -204,6 +204,9 @@ data class Assignment(val to: String, val expression: Expression) : Code() {
 }
 
 data class GroupedCode(val list: List<Code>) : Code() {
+
+    constructor(vararg code: Code) : this(code.toList())
+
     override fun toCodeBlock(): CustomCodeBlock {
         val blocks = list.map { it.toCodeBlock() }
 

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

@@ -1,46 +1,69 @@
 package com.virtuslab.pulumikotlin.codegen.step2intermediate
 
+import com.pulumi.kotlin.ConvertibleToJava
 import com.squareup.kotlinpoet.ClassName
 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
 
 object MoreTypes {
 
     object Kotlin {
         object Pulumi {
-            fun toJava(): MemberName {
+            fun toJavaMethod(): MemberName {
                 return MemberName("com.pulumi.kotlin", "toJava")
             }
 
-            fun applySuspend(): MemberName {
+            fun pulumiDslMarkerAnnotation(): ClassName {
+                return ClassName("com.pulumi.kotlin", "PulumiTagMarker")
+            }
+
+            fun toKotlinMethod(): MemberName {
+                return MemberName("com.pulumi.kotlin", "toKotlin")
+            }
+
+            fun applySuspendExtensionMethod(): MemberName {
                 return MemberName("com.pulumi.kotlin", "applySuspend")
             }
 
-            fun ConvertibleToJava(type: TypeName): ParameterizedTypeName =
-                ClassName("com.pulumi.kotlin", "ConvertibleToJava").parameterizedBy(type)
+            fun applyValueExtensionMethod(): MemberName {
+                return MemberName("com.pulumi.kotlin", "applyValue")
+            }
+
+            fun convertibleToJavaClass(type: TypeName): ParameterizedTypeName =
+                ConvertibleToJava::class.asClassName().parameterizedBy(type)
+
+            fun customArgsClass(): ClassName {
+                return ClassName("com.pulumi.kotlin", "CustomArgs")
+            }
+
+            fun customArgsBuilderClass(): ClassName {
+                return ClassName("com.pulumi.kotlin", "CustomArgsBuilder")
+            }
         }
 
-        fun Pair(leftType: TypeName, rightType: TypeName): ParameterizedTypeName {
+        fun coroutinesFutureAwaitExtensionMethod(): MemberName {
+            return MemberName("kotlinx.coroutines.future", "await")
+        }
+
+        fun pairClass(leftType: TypeName, rightType: TypeName): ParameterizedTypeName {
             return ClassName("kotlin", "Pair").parameterizedBy(leftType, rightType)
         }
     }
 
     object Java {
         object Pulumi {
+            fun outputOfMethod(): MemberName =
+                outputClass().member("of")
 
-            object Output {
-                val of: MemberName
-                    get() = Output().member("of")
-            }
-
-            fun Output(): ClassName {
+            fun outputClass(): ClassName {
                 return ClassName("com.pulumi.core", "Output")
             }
 
-            fun Output(type: TypeName): ParameterizedTypeName {
+            fun outputClass(type: TypeName): ParameterizedTypeName {
                 return ClassName("com.pulumi.core", "Output").parameterizedBy(type)
             }
         }

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

@@ -2,7 +2,7 @@ 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
 
 sealed class FieldType<T : ReferencedType> {
@@ -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(type.toTypeName())
     }
 }
 
 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,39 @@
+package com.virtuslab.pulumikotlin.codegen.step3codegen
+
+import com.squareup.kotlinpoet.CodeBlock
+import com.squareup.kotlinpoet.FileSpec
+import com.squareup.kotlinpoet.FunSpec
+import com.squareup.kotlinpoet.MemberName
+import com.squareup.kotlinpoet.PropertySpec
+import com.squareup.kotlinpoet.TypeSpec
+import com.virtuslab.pulumikotlin.codegen.expressions.Code
+
+object KotlinPoetExtensions {
+    fun FileSpec.Builder.addTypes(vararg typeSpecs: TypeSpec) =
+        addTypes(typeSpecs.toList())
+
+    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): FileSpec.Builder {
+        memberNames.forEach {
+            addImport(it)
+        }
+        return this
+    }
+
+    fun FileSpec.Builder.addTypes(types: Iterable<TypeSpec>): FileSpec.Builder {
+        types.forEach { addType(it) }
+        return this
+    }
+
+    fun CodeBlock.Builder.add(code: Code): CodeBlock.Builder {
+        return this.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 ?: ""}")
+        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)

src/main/kotlin/com/virtuslab/pulumikotlin/codegen/step3codegen/resources/ResourceGenerator.kt

@@ -15,7 +15,10 @@ import com.virtuslab.pulumikotlin.codegen.step2intermediate.Depth.Root
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.Direction.Output
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.LanguageType.Java
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.LanguageType.Kotlin
-import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Java.Pulumi.outputClass
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Kotlin.Pulumi.customArgsBuilderClass
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Kotlin.Pulumi.customArgsClass
+import com.virtuslab.pulumikotlin.codegen.step2intermediate.MoreTypes.Kotlin.Pulumi.pulumiDslMarkerAnnotation
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.NamingFlags
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.ResourceType
 import com.virtuslab.pulumikotlin.codegen.step2intermediate.Subject.Resource
@@ -42,10 +45,6 @@ object ResourceGenerator {
     }
 
     private fun buildArgsClass(fileSpecBuilder: FileSpec.Builder, resourceType: ResourceType) {
-        val dslTag = ClassName("com.pulumi.kotlin", "PulumiTagMarker")
-
-        val customArgs = ClassName("com.pulumi.kotlin", "CustomArgs")
-
         val javaFlags = NamingFlags(Root, Resource, Output, Java)
         val kotlinFlags = NamingFlags(Root, Resource, Output, Kotlin)
 
@@ -56,7 +55,7 @@ object ResourceGenerator {
         val fields = resourceType.outputFields.map { field ->
             PropertySpec.builder(
                 field.name,
-                MoreTypes.Java.Pulumi.Output(
+                outputClass(
                     field.fieldType.type.toTypeName().copy(nullable = !field.required),
                 ),
             )
@@ -120,11 +119,11 @@ object ResourceGenerator {
             .addParameter(
                 "block",
                 LambdaTypeName.get(
-                    ClassName("com.pulumi.kotlin", "CustomArgsBuilder"),
+                    customArgsBuilderClass(),
                     returnType = UNIT,
                 ).copy(suspending = true),
             )
-            .addStatement("val builder = %T()", ClassName("com.pulumi.kotlin", "CustomArgsBuilder"))
+            .addStatement("val builder = %T()", customArgsBuilderClass())
             .addStatement("block(builder)")
             .addStatement("this.opts = builder.build()")
             .addDocs("@param block A bag of options that control this resource's behavior.")
@@ -137,7 +136,7 @@ object ResourceGenerator {
                     .addModifiers(INTERNAL)
                     .build(),
             )
-            .addAnnotation(dslTag)
+            .addAnnotation(pulumiDslMarkerAnnotation())
             .addProperties(
                 listOf(
                     PropertySpec.builder("name", STRING.copy(nullable = true))
@@ -148,9 +147,9 @@ object ResourceGenerator {
                         .mutable(true)
                         .initializer("null")
                         .build(),
-                    PropertySpec.builder("opts", customArgs)
+                    PropertySpec.builder("opts", customArgsClass())
                         .mutable(true)
-                        .initializer("%T()", customArgs)
+                        .initializer("%T()", customArgsClass())
                         .build(),
                 ),
             )