[FIR generator] Extract BaseTransformerTypeFinder into the common module

We want to use it in other tree generators.

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.tree.generator.model.Field
 import org.jetbrains.kotlin.fir.tree.generator.model.Implementation
 import org.jetbrains.kotlin.fir.tree.generator.model.LeafBuilder
 import org.jetbrains.kotlin.fir.tree.generator.printer.invisibleField
-import org.jetbrains.kotlin.fir.tree.generator.util.traverseParents
+import org.jetbrains.kotlin.generators.tree.traverseParents
 
 object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTreeBuilder) {
     fun configureBuilders() = with(firTreeBuilder) {

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/Main.kt

@@ -5,9 +5,11 @@
 
 package org.jetbrains.kotlin.fir.tree.generator
 
+import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.printer.generateElements
 import org.jetbrains.kotlin.fir.tree.generator.util.configureInterfacesAndAbstractClasses
-import org.jetbrains.kotlin.fir.tree.generator.util.detectBaseTransformerTypes
+import org.jetbrains.kotlin.generators.tree.Model
+import org.jetbrains.kotlin.generators.tree.detectBaseTransformerTypes
 import org.jetbrains.kotlin.generators.tree.addPureAbstractElement
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil.collectPreviouslyGeneratedFiles
@@ -20,7 +22,8 @@ fun main(args: Array<String>) {
         ?: File("../../tree/gen").canonicalFile
 
     NodeConfigurator.configureFields()
-    detectBaseTransformerTypes(FirTreeBuilder)
+    val model = Model(FirTreeBuilder.elements, AbstractFirTreeBuilder.baseFirElement)
+    detectBaseTransformerTypes(model)
     ImplementationConfigurator.configureImplementations()
     configureInterfacesAndAbstractClasses(FirTreeBuilder)
     addPureAbstractElement(FirTreeBuilder.elements, pureAbstractElementType)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Element.kt

@@ -76,9 +76,6 @@ class Element(name: String, override val propertyName: String, kind: Kind) : Abs
     override val transformableChildren: List<Field>
         get() = emptyList() // Use Implementation#transformableChildren instead
 
-    var baseTransformerType: Element? = null
-    val transformerClass: Element get() = baseTransformerType ?: this
-
     override val visitorParameterName: String
         get() = safeDecapitalizedName
 
@@ -162,10 +159,6 @@ class Element(name: String, override val propertyName: String, kind: Kind) : Abs
         result.values.toList()
     }
 
-    val allFirFields: List<Field> by lazy {
-        allFields.filter { it.isFirType }
-    }
-
     override fun toString(): String {
         return with(ImportCollector("")) { render() }
     }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt

@@ -6,21 +6,14 @@
 package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 
 sealed class Field : AbstractField() {
     open var withReplace: Boolean = false
-    abstract val isFirType: Boolean
 
     open var needsSeparateTransform: Boolean = false
     var parentHasSeparateTransform: Boolean = true
     open var needTransformInOtherChildren: Boolean = false
 
-    /**
-     * @see org.jetbrains.kotlin.fir.tree.generator.util.detectBaseTransformerTypes
-     */
-    var useInBaseTransformerDetection = true
-
     open var customInitializationCall: String? = null
 
     open val isMutableOrEmptyList: Boolean
@@ -79,7 +72,8 @@ class FieldWithDefault(val origin: Field) : Field() {
     override var withReplace: Boolean
         get() = origin.withReplace
         set(_) {}
-    override val isFirType: Boolean get() = origin.isFirType
+    override val containsElement: Boolean
+        get() = origin.containsElement
     override var needsSeparateTransform: Boolean
         get() = origin.needsSeparateTransform
         set(_) {}
@@ -149,7 +143,6 @@ class SimpleField(
     override var isLateinit: Boolean = false,
     override var isParameter: Boolean = false,
 ) : Field() {
-    override val isFirType: Boolean = false
     override var isMutable: Boolean = withReplace
 
     override fun internalCopy(): Field {
@@ -186,11 +179,10 @@ class FirField(
     override var withReplace: Boolean,
 ) : Field() {
 
-    override val typeRef: TypeRefWithNullability
+    override val typeRef: ElementRef
         get() = element
     override var isVolatile: Boolean = false
     override var isFinal: Boolean = false
-    override val isFirType: Boolean = true
 
     override var isMutable: Boolean = true
     override var isLateinit: Boolean = false
@@ -211,13 +203,17 @@ class FirField(
 
 class FieldList(
     override val name: String,
-    val baseType: TypeRef,
+    override val baseType: TypeRef,
     override var withReplace: Boolean,
     useMutableOrEmpty: Boolean = false
-) : Field() {
+) : Field(), ListField {
     override var defaultValueInImplementation: String? = null
+
     override val typeRef: ClassRef<PositionTypeParameterRef>
-        get() = StandardTypes.list.withArgs(baseType)
+        get() = super.typeRef
+
+    override val listType: ClassRef<PositionTypeParameterRef>
+        get() = StandardTypes.list
 
     override var isVolatile: Boolean = false
     override var isFinal: Boolean = false
@@ -234,6 +230,4 @@ class FieldList(
             isMutableOrEmptyList
         )
     }
-
-    override val isFirType: Boolean = baseType is GenericElementOrRef<*, *> && baseType.element is Element
 }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Implementation.kt

@@ -65,7 +65,7 @@ class Implementation(val element: Element, val name: String?) : FieldContainer,
         get() = true
 
     override val walkableChildren: List<FieldWithDefault>
-        get() = allFields.filter { it.isFirType && !it.withGetter && it.needAcceptAndTransform }
+        get() = allFields.filter { it.containsElement && !it.withGetter && it.needAcceptAndTransform }
 
     override val transformableChildren: List<FieldWithDefault>
         get() = walkableChildren.filter { it.isMutable }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/implementation.kt

@@ -285,7 +285,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                     continue
                 }
                 printBlock {
-                    if (field.isMutable && field.isFirType) {
+                    if (field.isMutable && field.containsElement) {
                         // TODO: replace with smth normal
                         if (typeName == "FirWhenExpressionImpl" && field.name == "subject") {
                             println(
@@ -314,7 +314,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                 } else {
                     printBlock {
                         for (field in allFields) {
-                            if (!field.isMutable || !field.isFirType || field.name == "subjectVariable") continue
+                            if (!field.isMutable || !field.containsElement || field.name == "subjectVariable") continue
                             if (!field.needsSeparateTransform) {
                                 field.transform()
                             }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/util/BaseTransformerTypeFinder.kt

@@ -1,42 +0,0 @@
-/*
- * Copyright 2010-2020 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.fir.tree.generator.util
-
-import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
-import org.jetbrains.kotlin.fir.tree.generator.model.Element
-import org.jetbrains.kotlin.fir.tree.generator.model.Field
-import org.jetbrains.kotlin.fir.tree.generator.model.FieldList
-import org.jetbrains.kotlin.fir.tree.generator.model.FirField
-import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
-
-/**
- * For each FIR element, sets its [Element.baseTransformerType] to one of it parents if that parent FIR type is used at least once as
- * a type of [Field], except when that [Field] is explicitly opted out of it via [Field.useInBaseTransformerDetection].
- */
-fun detectBaseTransformerTypes(builder: AbstractFirTreeBuilder) {
-    val usedAsFieldType = hashSetOf<Element>()
-    for (element in builder.elements) {
-        for (field in element.allFirFields) {
-            if (!field.useInBaseTransformerDetection) continue
-            val fieldElement = when (field) {
-                is FirField -> field.element.element
-                is FieldList -> (field.baseType as GenericElementOrRef<*, *>).element as Element
-                else -> error("Invalid field type: $field")
-            }
-            if (fieldElement == AbstractFirTreeBuilder.baseFirElement) continue
-            usedAsFieldType.add(fieldElement)
-        }
-    }
-
-    for (element in builder.elements) {
-        element.traverseParents {
-            if (it in usedAsFieldType) {
-                element.baseTransformerType = it
-                return@traverseParents
-            }
-        }
-    }
-}

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/util/util.kt

@@ -5,16 +5,10 @@
 
 package org.jetbrains.kotlin.fir.tree.generator.util
 
-import org.jetbrains.kotlin.fir.tree.generator.model.Element
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
 import org.jetbrains.kotlin.utils.SmartPrinter
 import java.io.File
 
-fun Element.traverseParents(block: (Element) -> Unit) {
-    block(this)
-    elementParents.forEach { it.element.traverseParents(block) }
-}
-
 operator fun <K, V, U> MutableMap<K, MutableMap<V, U>>.set(k1: K, k2: V, value: U) {
     this.putIfAbsent(k1, mutableMapOf())
     val map = getValue(k1)

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/config/AbstractTreeBuilder.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.ir.generator.config
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.ir.generator.model.*
 import org.jetbrains.kotlin.ir.generator.model.ElementRef
+import org.jetbrains.kotlin.ir.generator.model.ListField
 import org.jetbrains.kotlin.ir.generator.model.Model
 import org.jetbrains.kotlin.types.Variance
 import kotlin.properties.PropertyDelegateProvider
@@ -81,7 +82,7 @@ abstract class AbstractTreeBuilder {
         }
         return ListField(
             name = name,
-            elementType = elementType ?: InferredOverriddenType,
+            baseType = elementType ?: InferredOverriddenType,
             listType = listType,
             isNullable = nullable,
             mutable = mutability == ListField.Mutability.Var,

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Model.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.ir.generator.model
 
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.ListField as AbstractListField
 import org.jetbrains.kotlin.ir.generator.BASE_PACKAGE
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.topologicalSort
@@ -53,7 +54,6 @@ class Element(
         get() = emptyMap()
 
     override var parentInVisitor: Element? = null
-    var transformerReturnType: Element? = null
 
     var typeKind: TypeKind? = null
         set(value) {
@@ -197,16 +197,16 @@ class SingleField(
 
 class ListField(
     name: String,
-    var elementType: TypeRef,
+    override var baseType: TypeRef,
     private val isNullable: Boolean,
-    private val listType: ClassRef<PositionTypeParameterRef>,
+    override val listType: ClassRef<PositionTypeParameterRef>,
     mutable: Boolean,
     isChild: Boolean,
     override val transformable: Boolean,
-) : Field(name, mutable, isChild) {
+) : Field(name, mutable, isChild), AbstractListField {
 
-    override val typeRef: TypeRefWithNullability
+    override val typeRef: ClassRef<PositionTypeParameterRef>
-        get() = listType.withArgs(elementType).copy(isNullable)
+        get() = listType.withArgs(baseType).copy(isNullable)
 
     enum class Mutability {
         Immutable,

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/model/Transformations.kt

@@ -61,7 +61,7 @@ internal fun processFieldOverrides(elements: List<Element>) {
                                     transformInferredType(field.typeRef, (overriddenField as SingleField).typeRef) as TypeRefWithNullability
                             }
                             is ListField -> {
-                                field.elementType = transformInferredType(field.elementType, (overriddenField as ListField).elementType)
+                                field.baseType = transformInferredType(field.baseType, (overriddenField as ListField).baseType)
                             }
                         }
 

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/Elements.kt

@@ -12,6 +12,7 @@ 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.ir.generator.model.ListField
 import org.jetbrains.kotlin.ir.generator.model.Model
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
@@ -71,7 +72,7 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
                             is SingleField -> println(".accept(visitor, data)")
                             is ListField -> {
                                 print(".forEach { it")
-                                if (child.elementType.nullable) {
+                                if (child.baseType.nullable) {
                                     print("?")
                                 }
                                 println(".accept(visitor, data) }")

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/VisitorsAndTransformers.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.printer.*
 import org.jetbrains.kotlin.ir.generator.*
 import org.jetbrains.kotlin.ir.generator.model.*
+import org.jetbrains.kotlin.ir.generator.model.ListField
 import org.jetbrains.kotlin.ir.generator.model.Model
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.capitalizeAsciiOnly
 import org.jetbrains.kotlin.utils.SmartPrinter
@@ -270,7 +271,7 @@ private class TypeTransformerPrinter(
             .filter {
                 val type = when (it) {
                     is SingleField -> it.typeRef
-                    is ListField -> it.elementType
+                    is ListField -> it.baseType
                 }
                 type.toString() == irTypeType.toString()
             }

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractElement.kt

@@ -79,6 +79,25 @@ abstract class AbstractElement<Element, Field>(
 
     abstract override val transformableChildren: List<Field>
 
+    /**
+     * A custom return type of the corresponding transformer method for this element.
+     */
+    var transformerReturnType: Element? = null
+
+    /**
+     * @see org.jetbrains.kotlin.generators.tree.detectBaseTransformerTypes
+     */
+    internal var baseTransformerType: Element? = null
+
+    /**
+     * The return type of the corresponding transformer method for this element.
+     *
+     * By default, computed using [org.jetbrains.kotlin.generators.tree.detectBaseTransformerTypes], but can be customizaed via
+     * [transformerReturnType]
+     */
+    val transformerClass: Element
+        get() = transformerReturnType ?: baseTransformerType ?: element
+
     final override fun get(fieldName: String): Field? {
         return allFields.firstOrNull { it.name == fieldName }
     }

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractField.kt

@@ -38,8 +38,19 @@ abstract class AbstractField {
 
     var fromParent: Boolean = false
 
+    /**
+     * Whether this field can contain an element either directly or e.g. in a list.
+     */
+    open val containsElement: Boolean
+        get() = typeRef is ElementOrRef<*, *> || this is ListField && baseType is ElementOrRef<*, *>
+
     open val defaultValueInImplementation: String? get() = null
 
+    /**
+     * @see org.jetbrains.kotlin.generators.tree.detectBaseTransformerTypes
+     */
+    var useInBaseTransformerDetection = true
+
     override fun toString(): String {
         return name
     }

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/ListField.kt

@@ -0,0 +1,25 @@
+/*
+ * Copyright 2010-2023 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.generators.tree
+
+/**
+ * A field that is used to store a list with arbitrary elements.
+ */
+interface ListField {
+
+    /**
+     * The element type of the list.
+     */
+    val baseType: TypeRef
+
+    /**
+     * The list type of the field, e.g. [List] or [MutableList].
+     */
+    val listType: ClassRef<PositionTypeParameterRef>
+
+    val typeRef: ClassRef<PositionTypeParameterRef>
+        get() = listType.withArgs(baseType)
+}

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/utils.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2023 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.generators.tree
+
+fun <Element : AbstractElement<Element, *>> Element.traverseParents(block: (Element) -> Unit) {
+    block(this)
+    elementParents.forEach { it.element.traverseParents(block) }
+}
+
+/**
+ * For each tree element, sets its [AbstractElement.baseTransformerType] to one of it parents if that parent type is used at least once as
+ * a type of a field, except when that field is explicitly opted out of it via
+ * [AbstractField.useInBaseTransformerDetection].
+ */
+fun <Element : AbstractElement<Element, *>> detectBaseTransformerTypes(model: Model<Element>) {
+    val usedAsFieldType = hashSetOf<AbstractElement<*, *>>()
+    for (element in model.elements) {
+        for (field in element.allFields.filter { it.containsElement }) {
+            if (!field.useInBaseTransformerDetection) continue
+            val fieldElement = (field.typeRef as? ElementOrRef<*, *>)?.element
+                ?: ((field as? ListField)?.baseType as? ElementOrRef<*, *>)?.element
+                ?: continue
+            if (fieldElement.isRootElement) continue
+            usedAsFieldType.add(fieldElement)
+        }
+    }
+
+    for (element in model.elements) {
+        element.traverseParents {
+            if (it in usedAsFieldType) {
+                element.baseTransformerType = it
+                return@traverseParents
+            }
+        }
+    }
+}