[IR generator] Don't use kotlinpoet for auto-generating IR tree classes

KT-61970 Fixed KT-61703 Fixed
2023-10-27 18:55:48 +02:00
parent 63e9dd588c
commit ed28923282
52 changed files with 392 additions and 646 deletions
@@ -12,7 +12,6 @@ dependencies {
    implementation(project(":generators"))
    implementation(project(":generators:tree-generator-common"))
    implementation(project(":compiler:util"))
-    implementation("com.squareup:kotlinpoet:1.11.0")

    compileOnly(intellijCore())
    compileOnly(commonDependency("org.jetbrains.intellij.deps:trove4j"))
@@ -6,13 +6,14 @@
 package org.jetbrains.kotlin.ir.generator

 import org.jetbrains.kotlin.generators.tree.TypeKind
+import org.jetbrains.kotlin.generators.tree.type
 import org.jetbrains.kotlin.ir.generator.Packages.declarations
 import org.jetbrains.kotlin.ir.generator.Packages.exprs
 import org.jetbrains.kotlin.ir.generator.Packages.symbols
 import org.jetbrains.kotlin.ir.generator.Packages.tree
 import org.jetbrains.kotlin.ir.generator.Packages.types
+import org.jetbrains.kotlin.ir.generator.Packages.util
 import org.jetbrains.kotlin.ir.generator.Packages.visitors
-import org.jetbrains.kotlin.generators.tree.type

 object Packages {
    const val tree = "org.jetbrains.kotlin.ir"
@@ -22,6 +23,7 @@ object Packages {
    const val types = "org.jetbrains.kotlin.ir.types"
    const val visitors = "org.jetbrains.kotlin.ir.visitors"
    const val descriptors = "org.jetbrains.kotlin.descriptors"
+    const val util = "org.jetbrains.kotlin.ir.util"
 }

 val elementBaseType = type(tree, "IrElementBase", TypeKind.Class)
@@ -34,6 +36,7 @@ val mutableAnnotationContainerType = type(declarations, "IrMutableAnnotationCont
 val irTypeType = type(types, "IrType")
 val irFactoryType = type(declarations, "IrFactory")
 val stageControllerType = type(declarations, "StageController", TypeKind.Class)
+val idSignatureType = type(util, "IdSignature", TypeKind.Class)

 val symbolType = type(symbols, "IrSymbol")
 val packageFragmentSymbolType = type(symbols, "IrPackageFragmentSymbol")
@@ -57,3 +60,6 @@ val returnableBlockSymbolType = type(symbols, "IrReturnableBlockSymbol")
 val propertySymbolType = type(symbols, "IrPropertySymbol")
 val localDelegatedPropertySymbolType = type(symbols, "IrLocalDelegatedPropertySymbol")
 val typeAliasSymbolType = type(symbols, "IrTypeAliasSymbol")
+
+val obsoleteDescriptorBasedApiAnnotation = type(BASE_PACKAGE, "ObsoleteDescriptorBasedAPI", TypeKind.Class)
+val unsafeDuringIrConstructionApiAnnotation = type(symbols, "UnsafeDuringIrConstructionAPI", TypeKind.Class)
@@ -5,27 +5,26 @@

 package org.jetbrains.kotlin.ir.generator

-import com.squareup.kotlinpoet.*
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.SourceElement
 import org.jetbrains.kotlin.descriptors.ValueClassRepresentation
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.printer.FunctionParameter
+import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
 import org.jetbrains.kotlin.ir.generator.config.AbstractTreeBuilder
 import org.jetbrains.kotlin.ir.generator.config.ElementConfig
 import org.jetbrains.kotlin.ir.generator.config.ElementConfig.Category.*
+import org.jetbrains.kotlin.ir.generator.config.ListFieldConfig.Mutability.*
 import org.jetbrains.kotlin.ir.generator.config.ListFieldConfig.Mutability.Array
 import org.jetbrains.kotlin.ir.generator.config.ListFieldConfig.Mutability.List
-import org.jetbrains.kotlin.ir.generator.config.ListFieldConfig.Mutability.Var
 import org.jetbrains.kotlin.ir.generator.config.SimpleFieldConfig
 import org.jetbrains.kotlin.ir.generator.model.Element.Companion.elementName2typeName
-import org.jetbrains.kotlin.ir.generator.print.toPoet
-import org.jetbrains.kotlin.ir.generator.util.*
-import org.jetbrains.kotlin.ir.generator.util.Import
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.withIndent

 // Note the style of the DSL to describe IR elements, which is these things in the following order:
 // 1) config (see properties of ElementConfig)
@@ -38,10 +37,11 @@ object IrTree : AbstractTreeBuilder() {
    private fun descriptor(typeName: String, nullable: Boolean = false): SimpleFieldConfig =
        field(
            name = "descriptor",
-            type = ClassRef<TypeParameterRef>(TypeKind.Interface, Packages.descriptors, typeName),
+            type = type(Packages.descriptors, typeName),
            mutable = false,
            nullable = nullable,
        ) {
+            optInAnnotation = obsoleteDescriptorBasedApiAnnotation
            skipInIrFactory()
        }

@@ -55,16 +55,15 @@ object IrTree : AbstractTreeBuilder() {
            skipInIrFactory()
        }

-        val oldCallback = generationCallback
        generationCallback = {
-            oldCallback?.invoke(this)
-            addFunction(
-                FunSpec.builder("acquireSymbol")
-                    .addModifiers(KModifier.ABSTRACT)
-                    .addParameter("symbol", symbol.toPoet())
-                    .returns(this@element.toPoet())
-                    .build(),
+            println()
+            printFunctionDeclaration(
+                name = "acquireSymbol",
+                parameters = listOf(FunctionParameter("symbol", symbol)),
+                returnType = this@element,
+                modality = Modality.ABSTRACT,
            )
+            println()
        }
    }

@@ -208,7 +207,7 @@ object IrTree : AbstractTreeBuilder() {
        +field("isCrossinline", boolean)
        +field("isNoinline", boolean)
        +field("isHidden", boolean) {
-            additionalImports.add(Import("org.jetbrains.kotlin.ir.util", "IdSignature"))
+            usedTypes.add(idSignatureType)
            kDoc = """
            If `true`, the value parameter does not participate in [IdSignature] computation.

@@ -349,9 +348,7 @@ object IrTree : AbstractTreeBuilder() {
                   before IR for all sources is built (because fake-overrides of lazy classes may depend on
                   declaration of source classes, e.g. for java source classes)
            """.trimIndent()
-            generationCallback = {
-                addAnnotation(ClassName("org.jetbrains.kotlin.ir.symbols", "UnsafeDuringIrConstructionAPI"))
-            }
+            optInAnnotation = unsafeDuringIrConstructionApiAnnotation
        }
    }
    val typeParametersContainer: ElementConfig by element(Declaration) {
@@ -446,7 +443,7 @@ object IrTree : AbstractTreeBuilder() {
        fieldsToSkipInIrFactoryMethod.add("origin")

        +field("symbol", symbolType, mutable = false) {
-            baseGetter = code("error(\"Should never be called\")")
+            baseGetter = "error(\"Should never be called\")"
            skipInIrFactory()
        }
    }
@@ -497,16 +494,18 @@ object IrTree : AbstractTreeBuilder() {
        transformByChildren = true
        generateIrFactoryMethod = false

-        +descriptor("ModuleDescriptor")
+        +descriptor("ModuleDescriptor").apply {
+            optInAnnotation = null
+        }
        +field("name", type<Name>(), mutable = false)
        +field("irBuiltins", type(Packages.tree, "IrBuiltIns"), mutable = false)
        +listField("files", file, mutability = List, isChild = true)
-        val undefinedOffset = MemberName(Packages.tree, "UNDEFINED_OFFSET")
+        usedTypes += ArbitraryImportable(Packages.tree, "UNDEFINED_OFFSET")
        +field("startOffset", int, mutable = false) {
-            baseGetter = code("%M", undefinedOffset)
+            baseGetter = "UNDEFINED_OFFSET"
        }
        +field("endOffset", int, mutable = false) {
-            baseGetter = code("%M", undefinedOffset)
+            baseGetter = "UNDEFINED_OFFSET"
        }
    }
    val property: ElementConfig by element(Declaration) {
@@ -621,7 +620,9 @@ object IrTree : AbstractTreeBuilder() {
        parent(symbolOwner)

        +symbol(packageFragmentSymbolType)
-        +field("packageFragmentDescriptor", type(Packages.descriptors, "PackageFragmentDescriptor"), mutable = false)
+        +field("packageFragmentDescriptor", type(Packages.descriptors, "PackageFragmentDescriptor"), mutable = false) {
+            optInAnnotation = obsoleteDescriptorBasedApiAnnotation
+        }
        +field("moduleDescriptor", type(Packages.descriptors, "ModuleDescriptor"), mutable = false) {
            kDoc = """
            This should be a link to [IrModuleFragment] instead. 
@@ -632,16 +633,13 @@ object IrTree : AbstractTreeBuilder() {
        }
        +field("packageFqName", type<FqName>())
        +field("fqName", type<FqName>()) {
-            baseGetter = code("packageFqName")
-            generationCallback = {
-                val deprecatedAnnotation = AnnotationSpec.builder(Deprecated::class)
-                    .addMember(code("message = \"Please use `packageFqName` instead\""))
-                    .addMember(code("replaceWith = ReplaceWith(\"packageFqName\")"))
-                    .addMember(code("level = DeprecationLevel.ERROR"))
-                    .build()
-                addAnnotation(deprecatedAnnotation)
-                setter(FunSpec.setterBuilder().addParameter("value", FqName::class).addCode(code("packageFqName = value")).build())
-            }
+            baseGetter = "packageFqName"
+            customSetter = "packageFqName = value"
+            deprecation = Deprecated(
+                "Please use `packageFqName` instead",
+                ReplaceWith("packageFqName"),
+                DeprecationLevel.ERROR,
+            )
        }
    }
    val externalPackageFragment: ElementConfig by element(Declaration) {
@@ -679,11 +677,11 @@ object IrTree : AbstractTreeBuilder() {
        parent(attributeContainer)

        +field("attributeOwnerId", attributeContainer) {
-            baseDefaultValue = code("this")
+            baseDefaultValue = "this"
            skipInIrFactory()
        }
        +field("originalBeforeInline", attributeContainer, nullable = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
            skipInIrFactory()
        }
        +field("type", irTypeType)
@@ -740,67 +738,79 @@ object IrTree : AbstractTreeBuilder() {
        parent(declarationReference)

        +field("dispatchReceiver", expression, nullable = true, isChild = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
        +field("extensionReceiver", expression, nullable = true, isChild = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
        +symbol(s)
        +field("origin", statementOriginType, nullable = true)
        +listField("valueArguments", expression.copy(nullable = true), mutability = Array, isChild = true) {
-            generationCallback = {
-                addModifiers(KModifier.PROTECTED)
-            }
+            visibility = Visibility.PROTECTED
        }
        +listField("typeArguments", irTypeType.copy(nullable = true), mutability = Array) {
-            generationCallback = {
-                addModifiers(KModifier.PROTECTED)
-            }
+            visibility = Visibility.PROTECTED
        }

-        val checkArgumentSlotAccess = MemberName("org.jetbrains.kotlin.ir.expressions", "checkArgumentSlotAccess", true)
+        usedTypes += ArbitraryImportable(Packages.exprs, "checkArgumentSlotAccess")
        generationCallback = {
-            addFunction(
-                FunSpec.builder("getValueArgument")
-                    .addParameter("index", int.toPoet())
-                    .returns(expression.toPoet().copy(nullable = true))
-                    .addCode("%M(\"value\", index, valueArguments.size)\n", checkArgumentSlotAccess)
-                    .addCode("return valueArguments[index]")
-                    .build()
+            val indexParam = FunctionParameter("index", StandardTypes.int)
+            val valueArgumentParam = FunctionParameter("valueArgument", expression.copy(nullable = true))
+            val typeArgumentParam = FunctionParameter("type", irTypeType.copy(nullable = true))
+
+            fun printSizeProperty(listName: String) {
+                println()
+                println("val ", listName, "Count: Int")
+                withIndent {
+                    println("get() = ", listName, ".size")
+                }
+            }
+
+            printSizeProperty("valueArguments")
+            printSizeProperty("typeArguments")
+
+            fun printFunction(
+                name: String,
+                additionalParameter: FunctionParameter?,
+                returnType: TypeRefWithNullability,
+                vararg statements: String,
+            ) {
+                println()
+                printFunctionDeclaration(name, listOf(indexParam) + listOfNotNull(additionalParameter), returnType)
+                println(" {")
+                withIndent {
+                    statements.forEach { println(it) }
+                }
+                println("}")
+            }
+
+            printFunction(
+                "getValueArgument",
+                null,
+                expression.copy(nullable = true),
+                "checkArgumentSlotAccess(\"value\", index, valueArguments.size)",
+                "return valueArguments[index]",
            )
-            addFunction(
-                FunSpec.builder("getTypeArgument")
-                    .addParameter("index", int.toPoet())
-                    .returns(irTypeType.toPoet().copy(nullable = true))
-                    .addCode("%M(\"type\", index, typeArguments.size)\n", checkArgumentSlotAccess)
-                    .addCode("return typeArguments[index]")
-                    .build()
+            printFunction(
+                "getTypeArgument",
+                null,
+                irTypeType.copy(nullable = true),
+                "checkArgumentSlotAccess(\"type\", index, typeArguments.size)",
+                "return typeArguments[index]",
            )
-            addFunction(
-                FunSpec.builder("putValueArgument")
-                    .addParameter("index", int.toPoet())
-                    .addParameter("valueArgument", expression.toPoet().copy(nullable = true))
-                    .addCode("%M(\"value\", index, valueArguments.size)\n", checkArgumentSlotAccess)
-                    .addCode("valueArguments[index] = valueArgument")
-                    .build()
+            printFunction(
+                "putValueArgument",
+                valueArgumentParam,
+                StandardTypes.unit,
+                "checkArgumentSlotAccess(\"value\", index, valueArguments.size)",
+                "valueArguments[index] = valueArgument",
            )
-            addFunction(
-                FunSpec.builder("putTypeArgument")
-                    .addParameter("index", int.toPoet())
-                    .addParameter("type", irTypeType.toPoet().copy(nullable = true))
-                    .addCode("%M(\"type\", index, typeArguments.size)\n", checkArgumentSlotAccess)
-                    .addCode("typeArguments[index] = type")
-                    .build()
-            )
-            addProperty(
-                PropertySpec.builder("valueArgumentsCount", int.toPoet())
-                    .getter(FunSpec.getterBuilder().addCode("return valueArguments.size").build())
-                    .build()
-            )
-            addProperty(
-                PropertySpec.builder("typeArgumentsCount", int.toPoet())
-                    .getter(FunSpec.getterBuilder().addCode("return typeArguments.size").build())
-                    .build()
+            printFunction(
+                "putTypeArgument",
+                typeArgumentParam,
+                StandardTypes.unit,
+                "checkArgumentSlotAccess(\"type\", index, typeArguments.size)",
+                "typeArguments[index] = type",
            )
        }
    }
@@ -867,10 +877,7 @@ object IrTree : AbstractTreeBuilder() {

        +field("origin", statementOriginType, nullable = true)
        +listField("statements", statement, mutability = List, isChild = true) {
-            generationCallback = {
-                addModifiers(KModifier.OVERRIDE)
-            }
-            baseDefaultValue = code("ArrayList(2)")
+            baseDefaultValue = "ArrayList(2)"
        }
    }
    val block: ElementConfig by element(Expression) {
@@ -913,7 +920,7 @@ object IrTree : AbstractTreeBuilder() {

        +field("loop", loop)
        +field("label", string, nullable = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
    }
    val `break` by element(Expression) {
@@ -993,19 +1000,22 @@ object IrTree : AbstractTreeBuilder() {
        parent(expression)

        generationCallback = {
-            addFunction(
-                FunSpec.builder("contentEquals")
-                    .addModifiers(KModifier.ABSTRACT)
-                    .addParameter("other", constantValue.toPoet())
-                    .returns(boolean.toPoet())
-                    .build()
+            println()
+            printFunctionDeclaration(
+                name = "contentEquals",
+                parameters = listOf(FunctionParameter("other", constantValue)),
+                returnType = StandardTypes.boolean,
+                modality = Modality.ABSTRACT,
            )
-            addFunction(
-                FunSpec.builder("contentHashCode")
-                    .addModifiers(KModifier.ABSTRACT)
-                    .returns(int.toPoet())
-                    .build()
+            println()
+            println()
+            printFunctionDeclaration(
+                name = "contentHashCode",
+                parameters = emptyList(),
+                returnType = StandardTypes.int,
+                modality = Modality.ABSTRACT,
            )
+            println()
        }
    }
    val constantPrimitive: ElementConfig by element(Expression) {
@@ -1093,7 +1103,7 @@ object IrTree : AbstractTreeBuilder() {
        +symbol(fieldSymbolType, mutable = true)
        +field("superQualifierSymbol", classSymbolType, nullable = true)
        +field("receiver", expression, nullable = true, isChild = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
        +field("origin", statementOriginType, nullable = true)
    }
@@ -1141,11 +1151,11 @@ object IrTree : AbstractTreeBuilder() {

        +field("origin", statementOriginType, nullable = true)
        +field("body", expression, nullable = true, isChild = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
        +field("condition", expression, isChild = true)
        +field("label", string, nullable = true) {
-            baseDefaultValue = code("null")
+            baseDefaultValue = "null"
        }
    }
    val whileLoop: ElementConfig by element(Expression) {
@@ -13,7 +13,8 @@ import org.jetbrains.kotlin.ir.generator.print.*
 import java.io.File

 const val BASE_PACKAGE = "org.jetbrains.kotlin.ir"
-const val VISITOR_PACKAGE = "$BASE_PACKAGE.visitors"
+
+internal const val TREE_GENERATOR_README = "compiler/ir/ir.tree/tree-generator/ReadMe.md"

 fun main(args: Array<String>) {
    val generationPath = args.firstOrNull()?.let { File(it) }
@@ -5,13 +5,10 @@

 package org.jetbrains.kotlin.ir.generator.config

-import com.squareup.kotlinpoet.CodeBlock
-import com.squareup.kotlinpoet.PropertySpec
-import com.squareup.kotlinpoet.TypeSpec
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.ir.generator.BASE_PACKAGE
 import org.jetbrains.kotlin.ir.generator.model.Element
-import org.jetbrains.kotlin.ir.generator.util.*
+import org.jetbrains.kotlin.utils.SmartPrinter

 class Config(
    val elements: List<ElementConfig>,
@@ -26,7 +23,7 @@ class ElementConfig(
    val params = mutableListOf<TypeVariable>()
    val parents = mutableListOf<TypeRef>()
    val fields = mutableListOf<FieldConfig>()
-    val additionalImports = mutableListOf<Import>()
+    val usedTypes = mutableListOf<Importable>()

    var visitorName: String? = null
    var visitorParent: ElementConfig? = null
@@ -53,7 +50,7 @@ class ElementConfig(

    var typeKind: TypeKind? = null

-    var generationCallback: (TypeSpec.Builder.() -> Unit)? = null
+    var generationCallback: (context(ImportCollector) SmartPrinter.() -> Unit)? = null
    var kDoc: String? = null

    override val element get() = this
@@ -139,10 +136,14 @@ sealed class FieldConfig(
    val name: String,
    val isChild: Boolean,
 ) {
-    var baseDefaultValue: CodeBlock? = null
-    var baseGetter: CodeBlock? = null
-    var printProperty = true
-    var strictCastInTransformChildren = false
+    var baseDefaultValue: String? = null
+    var baseGetter: String? = null
+    var customSetter: String? = null
+    var optInAnnotation: ClassRef<*>? = null
+
+    var deprecation: Deprecated? = null
+
+    var visibility = Visibility.PUBLIC

    internal var useFieldInIrFactoryStrategy: UseFieldAsParameterInIrFactoryStrategy =
        if (isChild) UseFieldAsParameterInIrFactoryStrategy.No else UseFieldAsParameterInIrFactoryStrategy.Yes(null)
@@ -157,8 +158,6 @@ sealed class FieldConfig(

    var kDoc: String? = null

-    var generationCallback: (PropertySpec.Builder.() -> Unit)? = null
-
    override fun toString() = name
 }

@@ -5,11 +5,9 @@

 package org.jetbrains.kotlin.ir.generator.model

-import com.squareup.kotlinpoet.CodeBlock
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.ir.generator.config.ElementConfig
 import org.jetbrains.kotlin.ir.generator.config.FieldConfig
-import org.jetbrains.kotlin.ir.generator.util.*
 import org.jetbrains.kotlin.utils.topologicalSort
 import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef
@@ -77,7 +75,7 @@ class Element(

    override val kDoc = config.kDoc

-    val additionalImports: List<Import> = config.additionalImports
+    val usedTypes: List<Importable> = config.usedTypes

    override fun toString() = name

@@ -112,25 +110,33 @@ class Element(
 typealias ElementRef = GenericElementRef<Element, Field>
 typealias ElementOrRef = GenericElementOrRef<Element, Field>

+@Suppress("LeakingThis")
 sealed class Field(
    config: FieldConfig,
    override val name: String,
    override var isMutable: Boolean,
    val isChild: Boolean,
 ) : AbstractField() {
-    abstract val baseDefaultValue: CodeBlock?
-    abstract val baseGetter: CodeBlock?
-    var isOverride = false
-    var needsDescriptorApiAnnotation = false
+    abstract val baseDefaultValue: String?
+    abstract val baseGetter: String?
    abstract val transformable: Boolean

    val useInIrFactoryStrategy = config.useFieldInIrFactoryStrategy
+
    init {
        kDoc = config.kDoc
+        optInAnnotation = config.optInAnnotation
+        deprecation = config.deprecation
+        visibility = config.visibility
    }

-    val printProperty = config.printProperty
-    val generationCallback = config.generationCallback
+    override val withGetter: Boolean
+        get() = baseGetter != null
+
+    override val defaultValueInImplementation: String?
+        get() = baseGetter ?: baseDefaultValue
+
+    override val customSetter: String? = config.customSetter

    override fun toString() = "$name: $typeRef"

@@ -138,7 +144,7 @@ sealed class Field(
        get() = false

    override val isFinal: Boolean
-        get() = false
+        get() = defaultValueInImplementation != null

    override val isLateinit: Boolean
        get() = false
@@ -153,8 +159,8 @@ class SingleField(
    override var typeRef: TypeRefWithNullability,
    mutable: Boolean,
    isChild: Boolean,
-    override val baseDefaultValue: CodeBlock?,
-    override val baseGetter: CodeBlock?,
+    override val baseDefaultValue: String?,
+    override val baseGetter: String?,
 ) : Field(config, name, mutable, isChild) {
    override val transformable: Boolean
        get() = isMutable
@@ -168,8 +174,8 @@ class ListField(
    mutable: Boolean,
    isChild: Boolean,
    override val transformable: Boolean,
-    override val baseDefaultValue: CodeBlock?,
-    override val baseGetter: CodeBlock?,
+    override val baseDefaultValue: String?,
+    override val baseGetter: String?,
 ) : Field(config, name, mutable, isChild) {
    override val typeRef: TypeRefWithNullability
        get() = listType.withArgs(elementType)
@@ -50,7 +50,6 @@ fun config2model(config: Config): Model {
    configureInterfacesAndAbstractClasses(elements)
    addPureAbstractElement(elements, elementBaseType)
    markLeaves(elements)
-    configureDescriptorApiAnnotation(elements)
    processFieldOverrides(elements)
    addWalkableChildren(elements)

@@ -160,19 +159,6 @@ private fun markLeaves(elements: List<Element>) {
    }
 }

-private fun configureDescriptorApiAnnotation(elements: List<Element>) {
-    for (el in elements) {
-        for (field in el.fields) {
-            val type = field.typeRef
-            if (type is ClassRef<*> && type.packageName.startsWith("org.jetbrains.kotlin.descriptors") &&
-                type.simpleName.endsWith("Descriptor") && type.simpleName != "ModuleDescriptor"
-            ) {
-                field.needsDescriptorApiAnnotation = true
-            }
-        }
-    }
-}
-
 private fun processFieldOverrides(elements: List<Element>) {
    for (element in iterateElementsParentFirst(elements)) {
        for (field in element.fields) {
@@ -180,9 +166,8 @@ private fun processFieldOverrides(elements: List<Element>) {
                for (parent in visited.elementParents) {
                    val overriddenField = parent.element.fields.singleOrNull { it.name == field.name }
                    if (overriddenField != null) {
-                        field.isOverride = true
-                        field.needsDescriptorApiAnnotation =
-                            field.needsDescriptorApiAnnotation || overriddenField.needsDescriptorApiAnnotation
+                        field.fromParent = true
+                        field.optInAnnotation = field.optInAnnotation ?: overriddenField.optInAnnotation

                        fun transformInferredType(type: TypeRef, overriddenType: TypeRef) =
                            type.takeUnless { it is InferredOverriddenType } ?: overriddenType
@@ -1,48 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.ir.generator.print
-
-import com.squareup.kotlinpoet.FileSpec
-import com.squareup.kotlinpoet.TypeSpec
-import org.jetbrains.kotlin.generators.tree.printer.GeneratedFile
-import org.jetbrains.kotlin.generators.tree.printer.printGeneratedType
-import org.jetbrains.kotlin.ir.generator.util.Import
-import java.io.File
-
-internal const val TREE_GENERATOR_README = "compiler/ir/ir.tree/tree-generator/ReadMe.md"
-
-fun printTypeCommon(
-    generationPath: File,
-    packageName: String,
-    type: TypeSpec,
-    additionalImports: List<Import> = emptyList(),
-): GeneratedFile = printGeneratedType(generationPath, TREE_GENERATOR_README, packageName, type.name!!) {
-    val code = FileSpec.builder(packageName, type.name!!)
-        .apply {
-            additionalImports.forEach { addImport(it.packageName, it.className) }
-        }
-        .indent("    ")
-        .addType(type)
-        .build()
-        .toString()
-        .replace("package $packageName\n\n", "")
-        .unbacktickIdentifiers("data", "value", "operator", "constructor", "delegate", "receiver", "field")
-        .replace("public ", "")
-        .replace(":\\s*Unit".toRegex(), "")
-        .replace("import kotlin\\..*\\n".toRegex(), "")
-        // Half-baked attempt to remove double indent generated by KotlinPoet, which is not idiomatic according to the Kotlin style guide
-        .replace("            visitor.visit", "        visitor.visit")
-        .replace("            accept(transformer, data)", "        accept(transformer, data)")
-    print(code)
-}
-
-private fun String.unbacktickIdentifiers(vararg identifiers: String): String {
-    var result = this
-    for (identifier in identifiers) {
-        result = result.replace("`$identifier`", identifier)
-    }
-    return result
-}
@@ -5,278 +5,172 @@

 package org.jetbrains.kotlin.ir.generator.print

-import com.squareup.kotlinpoet.*
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
-import org.jetbrains.kotlin.generators.tree.TypeKind
+import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.printer.*
 import org.jetbrains.kotlin.ir.generator.BASE_PACKAGE
+import org.jetbrains.kotlin.ir.generator.TREE_GENERATOR_README
 import org.jetbrains.kotlin.ir.generator.elementTransformerType
 import org.jetbrains.kotlin.ir.generator.elementVisitorType
 import org.jetbrains.kotlin.ir.generator.model.*
-import org.jetbrains.kotlin.generators.tree.TypeRefWithNullability
-import org.jetbrains.kotlin.generators.tree.typeKind
-import org.jetbrains.kotlin.generators.tree.printer.extendedKDoc
-import org.jetbrains.kotlin.ir.generator.util.tryParameterizedBy
+import org.jetbrains.kotlin.utils.SmartPrinter
+import org.jetbrains.kotlin.utils.withIndent
 import java.io.File
 import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef

-fun printElements(generationPath: File, model: Model) = sequence {
-    for (element in model.elements) {
+private val transformIfNeeded = ArbitraryImportable("$BASE_PACKAGE.util", "transformIfNeeded")
+private val transformInPlace = ArbitraryImportable("$BASE_PACKAGE.util", "transformInPlace")

-        val elementName = element.toPoet()
-        val selfParametrizedElementName = element.toPoetSelfParameterized()
+private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Element, Field>(printer) {

-        val elementType = when (element.kind?.typeKind) {
-            null -> error("Element's category not configured")
-            TypeKind.Class -> TypeSpec.classBuilder(elementName)
-            TypeKind.Interface -> TypeSpec.interfaceBuilder(elementName)
-        }.apply {
-            addModifiers(
-                when (element.kind) {
-                    ImplementationKind.SealedClass -> listOf(KModifier.SEALED)
-                    ImplementationKind.SealedInterface -> listOf(KModifier.SEALED)
-                    ImplementationKind.AbstractClass -> listOf(KModifier.ABSTRACT)
-                    ImplementationKind.FinalClass -> listOf(KModifier.FINAL)
-                    ImplementationKind.OpenClass -> listOf(KModifier.OPEN)
-                    else -> emptyList()
-                }
+    override fun makeFieldPrinter(printer: SmartPrinter) = object : AbstractFieldPrinter<Field>(printer) {
+        override fun forceMutable(field: Field) = field.isMutable
+    }
+
+    override fun defaultElementKDoc(element: Element) =
+        "A ${if (element.isLeaf) "leaf" else "non-leaf"} IR tree element."
+
+    override val separateFieldsWithBlankLine: Boolean
+        get() = true
+
+    context(ImportCollector)
+    override fun SmartPrinter.printAdditionalMethods(element: Element) {
+        element.generationCallback?.invoke(this@ImportCollector, this)
+
+        if (element.hasAcceptMethod) {
+            printAcceptMethod(
+                element = element,
+                visitorClass = elementVisitorType,
+                hasImplementation = !element.isRootElement,
+                kDoc = """
+                Runs the provided [visitor] on the IR subtree with the root at this node.
+                
+                @param visitor The visitor to accept.
+                @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
+                @return The value returned by the topmost `visit*` invocation.
+                """.trimIndent().takeIf { element.isRootElement }
            )
-            addTypeVariables(element.params.map { it.toPoet() })
+        }

-            val (classes, interfaces) = element.parentRefs.partition { it.typeKind == TypeKind.Class }
-            classes.singleOrNull()?.let {
-                superclass(it.toPoet())
-            }
-            addSuperinterfaces(interfaces.map { it.toPoet() })
+        if (element.hasTransformMethod) {
+            printTransformMethod(
+                element = element,
+                transformerClass = elementTransformerType,
+                implementation = "accept(transformer, data)".takeIf { !element.isRootElement },
+                returnType = element,
+                kDoc = """
+                    Runs the provided [transformer] on the IR subtree with the root at this node.
+        
+                    @param transformer The transformer to use.
+                    @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
+                    @return The transformed node.
+                    """.trimIndent().takeIf { element.isRootElement }
+            )
+        }

-            for (field in element.fields) {
-                if (!field.printProperty) continue
-                val poetType = field.typeRef.toPoet().copy(nullable = field.nullable)
-                addProperty(PropertySpec.builder(field.name, poetType).apply {
-                    mutable(field.isMutable)
-                    if (field.isOverride) {
-                        addModifiers(KModifier.OVERRIDE)
-                    }
-                    if (field.baseDefaultValue == null && field.baseGetter == null) {
-                        addModifiers(KModifier.ABSTRACT)
-                    }
-
-                    field.baseDefaultValue?.let {
-                        initializer(it)
-                    }
-                    field.baseGetter?.let {
-                        getter(FunSpec.getterBuilder().addCode("return ").addCode(it).build())
-                    }
-
-                    if (field.needsDescriptorApiAnnotation) {
-                        addAnnotation(descriptorApiAnnotation)
-                    }
-
-                    field.kDoc?.let {
-                        addKdoc(it)
-                    }
-
-                    field.generationCallback?.invoke(this)
-                }.build())
-            }
-
-            val isRootElement = element.isRootElement
-            val acceptMethodName = "accept"
-            val transformMethodName = "transform"
-            if (element.hasAcceptMethod) {
-                addFunction(FunSpec.builder(acceptMethodName).apply {
-                    addModifiers(if (isRootElement) KModifier.ABSTRACT else KModifier.OVERRIDE)
-                    val r = TypeVariableName("R")
-                    val d = TypeVariableName("D")
-                    addTypeVariable(r)
-                    addTypeVariable(d)
-                    val visitorParam = ParameterSpec.builder("visitor", elementVisitorType.toPoet().tryParameterizedBy(r, d))
-                        .build()
-                        .also(::addParameter)
-                    val dataParam = ParameterSpec.builder("data", d)
-                        .build()
-                        .also(::addParameter)
-                    returns(r)
-                    if (!isRootElement) {
-                        addStatement("return %N.%N(this, %N)", visitorParam, element.visitFunctionName, dataParam)
-                    }
-                    if (isRootElement) {
-                        addKdoc(
-                            """
-                            Runs the provided [%1N] on the IR subtree with the root at this node.
-
-                            @param %1N The visitor to accept.
-                            @param %2N An arbitrary context to pass to each invocation of [%1N]'s methods.
-                            @return The value returned by the topmost `visit*` invocation.
-                            """.trimIndent(),
-                            visitorParam,
-                            dataParam,
-                        )
-                    }
-                }.build())
-            }
-
-            if (element.hasTransformMethod) {
-                addFunction(FunSpec.builder(transformMethodName).apply {
-                    addModifiers(if (isRootElement) KModifier.ABSTRACT else KModifier.OVERRIDE)
-                    val d = TypeVariableName("D")
-                    addTypeVariable(d)
-                    val transformerParam = ParameterSpec.builder("transformer", elementTransformerType.toPoet().tryParameterizedBy(d))
-                        .build()
-                        .also(::addParameter)
-                    val dataParam = ParameterSpec.builder("data", d)
-                        .build()
-                        .also(::addParameter)
-                    returns(selfParametrizedElementName)
-                    if (!isRootElement) {
-                        addStatement("return %N(%N, %N) as %T", acceptMethodName, transformerParam, dataParam, selfParametrizedElementName)
-                    }
-                    if (isRootElement) {
-                        addKdoc(
-                            """
-                            Runs the provided [%1N] on the IR subtree with the root at this node.
-
-                            @param %1N The transformer to use.
-                            @param %2N An arbitrary context to pass to each invocation of [%1N]'s methods.
-                            @return The transformed node.
-                            """.trimIndent(),
-                            transformerParam,
-                            dataParam,
-                        )
-                    }
-                }.build())
-            }
-
-            if (element.hasAcceptChildrenMethod) {
-                addFunction(FunSpec.builder("acceptChildren").apply {
-                    addModifiers(if (isRootElement) KModifier.ABSTRACT else KModifier.OVERRIDE)
-                    val d = TypeVariableName("D")
-                    addTypeVariable(d)
-
-                    val visitorParam = ParameterSpec
-                        .builder("visitor", elementVisitorType.toPoet().tryParameterizedBy(UNIT, d)).build()
-                        .also(::addParameter)
-                    val dataParam = ParameterSpec
-                        .builder("data", d).build()
-                        .also(::addParameter)
+        if (element.hasAcceptChildrenMethod) {
+            printAcceptChildrenMethod(
+                element = element,
+                visitorClass = elementVisitorType,
+                visitorResultType = StandardTypes.unit,
+                override = !element.isRootElement,
+                kDoc = """
+                    Runs the provided [visitor] on subtrees with roots in this node's children.
+                    
+                    Basically, calls `accept(visitor, data)` on each child of this node.
+                    
+                    Does **not** run [visitor] on this node itself.
+                    
+                    @param visitor The visitor for children to accept.
+                    @param data An arbitrary context to pass to each invocation of [visitor]'s methods.
+                    """.trimIndent().takeIf { element.isRootElement }
+            )

+            if (!element.isRootElement) {
+                println(" {")
+                withIndent {
                    for (child in element.walkableChildren) {
-                        addStatement(buildString {
-                            append("%N")
-                            if (child.nullable) append("?")
-                            when (child) {
-                                is SingleField -> append(".%N(%N, %N)")
-                                is ListField -> {
-                                    append(".forEach { it")
-                                    if ((child.elementType as? TypeRefWithNullability)?.nullable == true) append("?")
-                                    append(".%N(%N, %N) }")
+                        print(child.name)
+                        if (child.nullable) {
+                            print("?")
+                        }
+                        when (child) {
+                            is SingleField -> println(".accept(visitor, data)")
+                            is ListField -> {
+                                print(".forEach { it")
+                                if (child.elementType.nullable) {
+                                    print("?")
                                }
+                                println(".accept(visitor, data) }")
                            }
-                        }, child.name, acceptMethodName, visitorParam, dataParam)
+                        }
                    }
-
-                    if (isRootElement) {
-                        addKdoc(
-                            """
-                            Runs the provided [%1N] on subtrees with roots in this node's children.
-                            
-                            Basically, calls `%3N(%1N, %2N)` on each child of this node.
-                            
-                            Does **not** run [%1N] on this node itself.
-                            
-                            @param %1N The visitor for children to accept.
-                            @param %2N An arbitrary context to pass to each invocation of [%1N]'s methods.
-                            """.trimIndent(),
-                            visitorParam,
-                            dataParam,
-                            acceptMethodName,
-                        )
-                    }
-                }.build())
+                }
+                print("}")
            }
+            println()
+        }

-            if (element.hasTransformChildrenMethod) {
-                addFunction(FunSpec.builder("transformChildren").apply {
-                    addModifiers(if (isRootElement) KModifier.ABSTRACT else KModifier.OVERRIDE)
-                    val d = TypeVariableName("D")
-                    addTypeVariable(d)
-                    val transformerParam =
-                        ParameterSpec.builder("transformer", elementTransformerType.toPoet().tryParameterizedBy(d)).build()
-                            .also(::addParameter)
-                    val dataParam = ParameterSpec.builder("data", d).build().also(::addParameter)
-
+        if (element.hasTransformChildrenMethod) {
+            printTransformChildrenMethod(
+                element = element,
+                transformerClass = elementTransformerType,
+                returnType = StandardTypes.unit,
+                override = !element.isRootElement,
+                kDoc = """
+                    Recursively transforms this node's children *in place* using [transformer].
+                    
+                    Basically, executes `this.child = this.child.transform(transformer, data)` for each child of this node.
+                    
+                    Does **not** run [transformer] on this node itself.
+                    
+                    @param transformer The transformer to use for transforming the children.
+                    @param data An arbitrary context to pass to each invocation of [transformer]'s methods.
+                    """.trimIndent().takeIf { element.isRootElement }
+            )
+            if (!element.isRootElement) {
+                println(" {")
+                withIndent {
                    for (child in element.transformableChildren) {
-                        val args = mutableListOf<Any>()
-                        val code = buildString {
-                            append("%N")
-                            args.add(child.name)
-                            when (child) {
-                                is SingleField -> {
-                                    append(" = %N")
-                                    args.add(child.name)
-                                    if (child.nullable) append("?")
-                                    append(".%N(%N, %N)")
-                                    args.add(transformMethodName)
+                        print(child.name)
+                        when (child) {
+                            is SingleField -> {
+                                print(" = ", child.name)
+                                if (child.nullable) {
+                                    print("?")
                                }
-                                is ListField -> {
-                                    if (child.isMutable) {
-                                        append(" = ")
-                                        append(child.name)
-                                        if (child.nullable) append("?")
-                                    }
-                                    append(".%M(%N, %N)")
-                                    args.add(if (child.isMutable) transformIfNeeded else transformInPlace)
+                                print(".transform(transformer, data)")
+                                val elementRef = child.typeRef as GenericElementRef<*, *>
+                                if (!elementRef.element.hasTransformMethod) {
+                                    print(" as ", elementRef.render())
                                }
+                                println()
                            }
-
-                            args.add(transformerParam)
-                            args.add(dataParam)
-
-                            if (child is SingleField) {
-                                @Suppress("UNCHECKED_CAST")
-                                val elRef = child.typeRef as GenericElementRef<Element, Field>
-                                if (!elRef.element.hasTransformMethod) {
-                                    append(" as %T")
-                                    args.add(elRef.toPoet())
+                            is ListField -> {
+                                if (child.isMutable) {
+                                    print(" = ", child.name)
+                                    if (child.nullable) {
+                                        print("?")
+                                    }
+                                    addImport(transformIfNeeded)
+                                    println(".transformIfNeeded(transformer, data)")
+                                } else {
+                                    addImport(transformInPlace)
+                                    println(".transformInPlace(transformer, data)")
                                }
                            }
                        }
-
-                        addStatement(code, *args.toTypedArray())
                    }
-
-                    if (isRootElement) {
-                        addKdoc(
-                            """
-                            Recursively transforms this node's children *in place* using [%1N].
-                            
-                            Basically, executes `this.child = this.child.%3N(%1N, %2N)` for each child of this node.
-                            
-                            Does **not** run [%1N] on this node itself.
-                            
-                            @param %1N The transformer to use for transforming the children.
-                            @param %2N An arbitrary context to pass to each invocation of [%1N]'s methods.
-                            """.trimIndent(),
-                            transformerParam,
-                            dataParam,
-                            transformMethodName,
-                        )
-                    }
-                }.build())
+                }
+                print("}")
            }
-
-            generateElementKDoc(element)
-
-            element.generationCallback?.invoke(this)
-        }.build()
-
-        yield(printTypeCommon(generationPath, elementName.packageName, elementType, element.additionalImports))
+            println()
+        }
    }
 }

-private fun TypeSpec.Builder.generateElementKDoc(element: Element) {
-    addKdoc(element.extendedKDoc("A ${if (element.isLeaf) "leaf" else "non-leaf"} IR tree element."))
+fun printElements(generationPath: File, model: Model) = model.elements.map { element ->
+    printGeneratedType(generationPath, TREE_GENERATOR_README, element.packageName, element.typeName) {
+        addAllImports(element.usedTypes)
+        ElementPrinter(this).printElement(element)
+    }
 }
-
-private val descriptorApiAnnotation = ClassName("org.jetbrains.kotlin.ir", "ObsoleteDescriptorBasedAPI")
-private val transformIfNeeded = MemberName("$BASE_PACKAGE.util", "transformIfNeeded", true)
-private val transformInPlace = MemberName("$BASE_PACKAGE.util", "transformInPlace", true)
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.generators.tree.printer.GeneratedFile
 import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
 import org.jetbrains.kotlin.generators.tree.printer.printGeneratedType
 import org.jetbrains.kotlin.ir.generator.IrTree
+import org.jetbrains.kotlin.ir.generator.TREE_GENERATOR_README
 import org.jetbrains.kotlin.ir.generator.config.UseFieldAsParameterInIrFactoryStrategy
 import org.jetbrains.kotlin.ir.generator.irFactoryType
 import org.jetbrains.kotlin.ir.generator.model.Element
@@ -28,7 +29,6 @@ internal fun printFactory(generationPath: File, model: Model): GeneratedFile = p
    irFactoryType.packageName,
    irFactoryType.simpleName,
 ) {
-    println()
    println("interface ", irFactoryType.simpleName, " {")
    withIndent {
        println("val stageController: ", stageControllerType.render())
@@ -1,68 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.ir.generator.print
-
-import com.squareup.kotlinpoet.*
-import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.ir.generator.config.ElementConfigOrRef
-import org.jetbrains.kotlin.ir.generator.model.Element
-import org.jetbrains.kotlin.ir.generator.model.Element.Companion.elementName2typeName
-import org.jetbrains.kotlin.ir.generator.util.*
-import org.jetbrains.kotlin.types.Variance
-import  org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef
-
-fun Element.toPoet() = ClassName(packageName, this.typeName)
-fun Element.toPoetSelfParameterized() = toPoet().parameterizedByIfAny(poetTypeVariables)
-fun Element.toPoetStarParameterized() = toPoet().parameterizedByIfAny(List(params.size) { STAR })
-val Element.poetTypeVariables get() = params.map { TypeVariableName(it.name) }
-
-fun TypeRef.toPoet(): TypeName {
-    return when (this) {
-        is GenericElementRef<*, *> -> ClassName(element.packageName, element.typeName).parameterizedByIfAny(typeArgsToPoet())
-        is ClassRef<*> -> ClassName(packageName, simpleNames).parameterizedByIfAny(typeArgsToPoet())
-        is NamedTypeParameterRef -> TypeVariableName(name)
-        is ElementConfigOrRef -> ClassName(this.element.category.packageName, elementName2typeName(element.name)).parameterizedByIfAny(
-            typeArgsToPoet()
-        ) // ad-hoc solution
-        is TypeRef.Star -> STAR
-        else -> error("Unexpected type reference: $this")
-    }.let {
-        if (this is TypeRefWithNullability) it.copy(nullable = nullable) else it
-    }
-}
-
-private fun ParametrizedTypeRef<*, *>.typeArgsToPoet(): List<TypeName> {
-    if (args.isEmpty()) {
-        return emptyList()
-    }
-
-    fun fromPositional(args: Map<PositionTypeParameterRef, TypeRef>): List<TypeName> {
-        val num = args.keys.maxOfOrNull { it.index }!!
-        return (0..num).map { i -> args[PositionTypeParameterRef(i)]?.toPoet() ?: STAR }
-    }
-
-    val positional = args.keys.filterIsInstance<PositionTypeParameterRef>()
-    if (positional.size == args.size) {
-        @Suppress("UNCHECKED_CAST")
-        return fromPositional(args as Map<PositionTypeParameterRef, TypeRef>)
-    } else {
-        check(positional.isEmpty()) { "Can't yet handle mixed index-name args" }
-        this as GenericElementRef<*, *> // Named args must only be used with generated elements (for now)
-
-        val args = args.entries
-            .associate { p -> PositionTypeParameterRef(element.params.withIndex().single { it.value.name == p.key.name }.index) to p.value }
-        return fromPositional(args)
-    }
-}
-
-fun TypeVariable.toPoet() = TypeVariableName(
-    name, bounds.map { it.toPoet() }, when (variance) {
-        Variance.INVARIANT -> null
-        Variance.IN_VARIANCE -> KModifier.IN
-        Variance.OUT_VARIANCE -> KModifier.OUT
-    }
-)
-
@@ -22,7 +22,6 @@ private fun printVisitorCommon(
    makePrinter: (SmartPrinter, ClassRef<*>) -> AbstractVisitorPrinter<Element, Field>,
 ): GeneratedFile =
    printGeneratedType(generationPath, TREE_GENERATOR_README, visitorType.packageName, visitorType.simpleName) {
-        println()
        makePrinter(this, visitorType).printVisitor(model.elements)
    }

@@ -1,29 +0,0 @@
-/*
- * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.ir.generator.util
-
-import com.squareup.kotlinpoet.ClassName
-import com.squareup.kotlinpoet.CodeBlock
-import com.squareup.kotlinpoet.ParameterizedTypeName
-import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
-import com.squareup.kotlinpoet.TypeName
-
-class Import(val packageName: String, val className: String)
-
-fun ClassName.parameterizedByIfAny(typeArguments: List<TypeName>) =
-    if (typeArguments.isNotEmpty()) parameterizedBy(typeArguments) else this
-
-fun TypeName.tryParameterizedBy(vararg typeArguments: TypeName) = when (this) {
-    is ClassName -> parameterizedBy(*typeArguments)
-    is ParameterizedTypeName -> this.rawType.parameterizedBy(*typeArguments)
-    else -> {
-        if (typeArguments.isNotEmpty())
-            error("Type $this cannot be parameterized")
-        else this
-    }
-}
-
-fun code(code: String, vararg args: Any?) = CodeBlock.of(code, *args)