[FIR generator] Use a single list of parent TypeRefs for elements

(Instead of storing the list of parent elements + the map of parent
arguments)

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

@@ -166,11 +166,11 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         loopJump.configure {
-            parentArg(jump, "E", loop)
+            parentArgs(jump, "E" to loop)
         }
 
         returnExpression.configure {
-            parentArg(jump, "E", function)
+            parentArgs(jump, "E" to function)
             +field("result", expression).withTransform()
             needTransformOtherChildren()
         }

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

@@ -34,19 +34,25 @@ abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val
             return TypeVariable(name, upperBounds, variance).also(element.params::add)
         }
 
-        fun parentArg(parent: Element, argument: String, type: TypeRef) {
+        fun parentArgs(parent: Element, vararg arguments: Pair<String, TypeRef>) {
-            parentArg(parent, NamedTypeParameterRef(argument), type)
+            parentArgs(parent, listOf(*arguments))
         }
 
-        fun parentArg(parent: Element, argument: NamedTypeParameterRef, type: TypeRef) {
+        private fun parentArgs(parent: Element, arguments: List<Pair<String, TypeRef>>) {
-            require(parent in element.parents) {
+            parentArgs(parent, arguments.map { (name, arg) -> NamedTypeParameterRef(name) to arg })
+        }
+
+        @JvmName("parentArgs2")
+        private fun parentArgs(parent: Element, arguments: List<Pair<NamedTypeParameterRef, TypeRef>>) {
+            val parentIndex = element.parentRefs.indexOfFirst { it.element == parent }
+            require(parentIndex >= 0) {
                 "$parent is not parent of $element"
             }
-            val argMap = element.parentsArguments.getOrPut(parent) { mutableMapOf() }
+            val parentRef = element.parentRefs[parentIndex]
-            require(argument !in argMap) {
+            require(parentRef.args.isEmpty()) {
-                "Argument $argument already defined for parent $parent of $element"
+                "Parent $parent of element $element already has type arguments: $parentRef"
             }
-            argMap[argument] = type
+            element.parentRefs[parentIndex] = parentRef.copy(arguments.toMap())
         }
 
         fun needTransformOtherChildren() {

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

@@ -47,9 +47,9 @@ abstract class AbstractFirTreeBuilder {
     private fun createElement(name: String, kind: Element.Kind, vararg dependencies: Element): Element =
         Element(name, kind).also {
             if (dependencies.isEmpty()) {
-                it.parents.add(baseFirElement)
+                it.parentRefs.add(baseFirElement)
             }
-            it.parents.addAll(dependencies)
+            it.parentRefs.addAll(dependencies)
             elements += it
         }
 

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

@@ -57,7 +57,7 @@ abstract class AbstractFirTreeImplementationConfigurator {
     private fun collectLeafsWithoutImplementation(builder: AbstractFirTreeBuilder): Set<Element> {
         val elements = builder.elements.toMutableSet()
         builder.elements.forEach {
-            elements.removeAll(it.parents)
+            elements.removeAll(it.parentRefs.map { it.element }.toSet())
         }
         elements.removeAll(elementsWithImpl)
         return elements

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.tree.generator.model
 import org.jetbrains.kotlin.fir.tree.generator.printer.BASE_PACKAGE
 import org.jetbrains.kotlin.fir.tree.generator.util.set
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.ElementOrRef
 import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef
 
@@ -38,13 +39,12 @@ class Element(override val name: String, kind: Kind) : AbstractElement<Element,
     override val type: String = "Fir$name"
     override val packageName: String = BASE_PACKAGE + kind.packageName.let { if (it.isBlank()) it else "." + it }
     override val fullQualifiedName: String get() = super.fullQualifiedName!!
-    override val parents = mutableListOf<Element>()
+    override val parentRefs = mutableListOf<ElementOrRef<Element, Field>>()
     var defaultImplementation: Implementation? = null
     val customImplementations = mutableListOf<Implementation>()
 
     override val params = mutableListOf<TypeVariable>()
 
-    override val parentsArguments = mutableMapOf<Element, MutableMap<TypeRef, TypeRef>>()
     override var kind: ImplementationKind? = null
         set(value) {
             if (value !in allowedKinds) {
@@ -61,7 +61,7 @@ class Element(override val name: String, kind: Kind) : AbstractElement<Element,
 
     var doesNotNeedImplementation: Boolean = false
 
-    val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || parents.any { it.needTransformOtherChildren }
+    val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || parentRefs.any { it.element.needTransformOtherChildren }
     override val overridenFields: MutableMap<Field, MutableMap<Field, Boolean>> = mutableMapOf()
     val useNullableForReplace: MutableSet<Field> = mutableSetOf()
     val allImplementations: List<Implementation> by lazy {
@@ -105,16 +105,18 @@ class Element(override val name: String, kind: Kind) : AbstractElement<Element,
 
     val parentFields: List<Field> by lazy {
         val result = LinkedHashMap<String, Field>()
-        parents.forEach { parent ->
+        parentRefs.forEach { parentRef ->
+            val parent = parentRef.element
             val fields = parent.allFields.map { field ->
                 val copy = (field as? SimpleField)?.let { simpleField ->
-                    parentsArguments[parent]?.get(NamedTypeParameterRef(simpleField.type))?.let {
+                    // FIXME: Replace with parentRef.args[simpleField.typeRef]
+                    parentRef.args[NamedTypeParameterRef(simpleField.type)]?.let {
                         simpleField.replaceType(it)
                     }
                 } ?: field.copy()
                 copy.apply {
                     arguments.replaceAll {
-                        parentsArguments[parent]?.get(it) ?: it
+                        parentRef.args[it] ?: it
                     }
                     fromParent = true
                 }

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

@@ -56,21 +56,22 @@ fun SmartPrinter.printElement(element: Element) {
         print(typeParameters())
         val needPureAbstractElement = !isInterface && !allParents.any { it.kind == ImplementationKind.AbstractClass || it.kind == ImplementationKind.SealedClass }
 
-        if (parents.isNotEmpty() || needPureAbstractElement) {
+        if (parentRefs.isNotEmpty() || needPureAbstractElement) {
             print(" : ")
             if (needPureAbstractElement) {
                 print("${pureAbstractElementType.type}()")
-                if (parents.isNotEmpty()) {
+                if (parentRefs.isNotEmpty()) {
                     print(", ")
                 }
             }
             print(
-                parents.joinToString(", ") {
+                parentRefs.joinToString(", ") {
-                    var result = it.type
+                    // TODO: Factor out
-                    parentsArguments[it]?.let { arguments ->
+                    var result = it.element.type
-                        result += arguments.values.joinToString(", ", "<", ">") { it.typeWithArguments }
+                    if (it.args.isNotEmpty()) {
+                        result += it.args.values.joinToString(", ", "<", ">") { it.typeWithArguments }
                     }
-                    result + it.kind.braces()
+                    result + it.element.kind.braces()
                 },
             )
         }
@@ -134,14 +135,16 @@ fun SmartPrinter.printElement(element: Element) {
             if (needTransformOtherChildren) {
                 println()
                 abstract()
-                if (element.parents.any { it.needTransformOtherChildren }) {
+                if (element.parentRefs.any { it.element.needTransformOtherChildren }) {
                     print("override ")
                 }
                 println(transformFunctionDeclaration("OtherChildren", typeWithArguments))
             }
 
             if (element == AbstractFirTreeBuilder.baseFirElement) {
-                require(isInterface)
+                require(isInterface) {
+                    "$element must be an interface"
+                }
                 println()
                 println("fun accept(visitor: FirVisitorVoid) = accept(visitor, null)")
                 println()

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

@@ -61,9 +61,9 @@ fun Implementation.collectImports(base: List<String> = emptyList(), kind: Import
 }
 
 fun Element.collectImports(): List<String> {
-    val baseTypes = parents.mapTo(mutableListOf()) { it.fullQualifiedName }
+    val baseTypes = parentRefs.mapTo(mutableListOf()) { it.fullQualifiedName!! }
     baseTypes += AbstractFirTreeBuilder.baseFirElement.fullQualifiedName
-    baseTypes += parentsArguments.values.flatMap { it.values }.mapNotNull { it.fullQualifiedName }
+    baseTypes += parentRefs.flatMap { it.args.values }.mapNotNull { it.fullQualifiedName } // TODO: Use recursive import collection for TypeRefs
     if (needPureAbstractElement) {
         baseTypes += pureAbstractElementType.fullQualifiedName
     }

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

@@ -21,12 +21,12 @@ private val elementsWithMultipleSupertypesForDefaultVisitor = mapOf(
 
 private fun Element.isAcceptableForDefaultVisiting(): Boolean {
     if (this == AbstractFirTreeBuilder.baseFirElement) return false
-    val hasSingleSupertype = parents.size == 1 && parents.single().name != "Element"
+    val hasSingleSupertype = parentRefs.size == 1 && parentRefs.single().element.name != "Element"
     return hasSingleSupertype || this in elementsWithMultipleSupertypesForDefaultVisitor
 }
 
 private fun Element.getNameOfSupertypeForDefaultVisiting(): String {
-    val parentForDefaultVisiting = parents.singleOrNull() ?: elementsWithMultipleSupertypesForDefaultVisitor.getValue(this)
+    val parentForDefaultVisiting = parentRefs.singleOrNull()?.element ?: elementsWithMultipleSupertypesForDefaultVisitor.getValue(this)
     return parentForDefaultVisiting.name
 }
 

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

@@ -12,7 +12,7 @@ import java.io.File
 
 fun Element.traverseParents(block: (Element) -> Unit) {
     block(this)
-    parents.forEach { it.traverseParents(block) }
+    parentRefs.forEach { it.element.traverseParents(block) }
 }
 
 operator fun <K, V, U> MutableMap<K, MutableMap<V, U>>.set(k1: K, k2: V, value: U) {

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

@@ -20,9 +20,7 @@ interface AbstractElement<Element, Field> : ElementOrRef<Element, Field>, FieldC
 
     val params: List<TypeVariable>
 
-    val parents: List<Element>
+    val parentRefs: List<ElementOrRef<Element, Field>>
-
-    val parentsArguments: Map<Element, Map<TypeRef, TypeRef>>
 
     val overridenFields: Map<Field, Map<Field, Boolean>>
 
@@ -30,7 +28,7 @@ interface AbstractElement<Element, Field> : ElementOrRef<Element, Field>, FieldC
         get() = false
 
     override val allParents: List<ImplementationKindOwner>
-        get() = parents
+        get() = parentRefs.map { it.element }
 
     override fun getTypeWithArguments(notNull: Boolean): String = type + generics
 }