[FIR generator] Replace AbstractElement with ElemenOrRef

In the FIR generator, the `AbstractElement` class was used to
represent either an element type without type arguments applied
(using the `Element` subclass), or an element type with applied type
arguments (using the `ElementWithArguments` subclass).

Instead, it is more logical to use the `Element` class to always
represent a non-parameterized element type, and for a parameterized
element type use the `ElementRef` class, just like we do
in the IR generator.

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

@@ -18,7 +18,6 @@ import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeProjection
 import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeRef
 import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.fir.tree.generator.model.*
-import org.jetbrains.kotlin.generators.tree.TypeRef
 
 object FieldSets {
     val calleeReference by lazy { field("calleeReference", reference, withReplace = true) }
@@ -35,7 +34,7 @@ object FieldSets {
 
     val arguments by lazy { fieldList("arguments", expression) }
 
-    val declarations by lazy { fieldList(declaration.withArgs("E" to TypeRef.Star)) }
+    val declarations by lazy { fieldList(declaration).apply { useInBaseTransformerDetection = false } }
 
     val annotations by lazy {
         fieldList(

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

@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFieldConfigurator
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder.Companion.baseFirElement
 import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.fir.tree.generator.model.*
-import org.jetbrains.kotlin.generators.tree.SimpleTypeArgument
 import org.jetbrains.kotlin.generators.tree.StandardTypes
 import org.jetbrains.kotlin.generators.tree.TypeRef
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
@@ -171,7 +170,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
 
         returnExpression.configure {
-            parentArg(jump, "E", function.withArgs("E" to TypeRef.Star))
+            parentArg(jump, "E", function)
             +field("result", expression).withTransform()
             needTransformOtherChildren()
         }
@@ -782,12 +781,3 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
         }
     }
 }
-
-// TODO: Replace with org.jetbrains.kotlin.generators.tree.withArgs
-fun Element.withArgs(vararg replacements: Pair<String, TypeRef>): AbstractElement {
-    val replaceMap = replacements.toMap()
-    val newArguments = params.map { param ->
-        replaceMap[param.name]?.let { SimpleTypeArgument(it.type, null) } ?: error("Type variable $param not found in $this")
-    }
-    return ElementWithArguments(this, newArguments)
-}

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

@@ -8,21 +8,10 @@ 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.AbstractElement as CommonAbstractElement
+import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
+import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef
 
-interface AbstractElement : CommonAbstractElement<AbstractElement, Field> {
+class Element(override val name: String, kind: Kind) : AbstractElement<Element, Field> {
-    val baseTransformerType: AbstractElement?
-    val transformerType: AbstractElement
-    val doesNotNeedImplementation: Boolean
-    val needTransformOtherChildren: Boolean
-    val allImplementations: List<Implementation>
-    val allFirFields: List<Field>
-    val defaultImplementation: Implementation?
-    val customImplementations: List<Implementation>
-    val useNullableForReplace: Set<Field>
-}
-
-class Element(override val name: String, kind: Kind) : AbstractElement {
     companion object {
         private val allowedKinds = setOf(
             ImplementationKind.Interface,
@@ -32,20 +21,30 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
         )
     }
 
+    override val element: Element
+        get() = this
+
+    override val args: Map<NamedTypeParameterRef, TypeRef>
+        get() = emptyMap()
+
+    override val nullable: Boolean
+        get() = false
+
+    override fun copy(nullable: Boolean) = ElementRef(this, args, nullable)
+
+    override fun copy(args: Map<NamedTypeParameterRef, TypeRef>) = ElementRef(this, args, nullable)
+
     override val fields = mutableSetOf<Field>()
     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 var defaultImplementation: Implementation? = null
+    var defaultImplementation: Implementation? = null
-    override val customImplementations = mutableListOf<Implementation>()
+    val customImplementations = mutableListOf<Implementation>()
 
     override val params = mutableListOf<TypeVariable>()
 
-    override val typeArguments: List<TypeArgument>
+    override val parentsArguments = mutableMapOf<Element, MutableMap<TypeRef, TypeRef>>()
-        get() = emptyList()
-
-    override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<TypeRef, TypeRef>>()
     override var kind: ImplementationKind? = null
         set(value) {
             if (value !in allowedKinds) {
@@ -57,15 +56,15 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
 
     override var isSealed: Boolean = false
 
-    override var baseTransformerType: Element? = null
+    var baseTransformerType: Element? = null
-    override val transformerType: Element get() = baseTransformerType ?: this
+    val transformerType: Element get() = baseTransformerType ?: this
 
-    override var doesNotNeedImplementation: Boolean = false
+    var doesNotNeedImplementation: Boolean = false
 
-    override val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || parents.any { it.needTransformOtherChildren }
+    val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || parents.any { it.needTransformOtherChildren }
     override val overridenFields: MutableMap<Field, MutableMap<Field, Boolean>> = mutableMapOf()
-    override val useNullableForReplace: MutableSet<Field> = mutableSetOf()
+    val useNullableForReplace: MutableSet<Field> = mutableSetOf()
-    override val allImplementations: List<Implementation> by lazy {
+    val allImplementations: List<Implementation> by lazy {
         if (doesNotNeedImplementation) {
             emptyList()
         } else {
@@ -139,7 +138,7 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
         result.values.toList()
     }
 
-    override val allFirFields: List<Field> by lazy {
+    val allFirFields: List<Field> by lazy {
         allFields.filter { it.isFirType }
     }
 
@@ -162,12 +161,6 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
     }
 }
 
-class ElementWithArguments(val element: Element, override val typeArguments: List<TypeArgument>) : AbstractElement by element {
+typealias ElementRef = GenericElementRef<Element, Field>
-    override fun equals(other: Any?): Boolean {
-        return element.equals(other)
-    }
 
-    override fun hashCode(): Int {
+typealias ElementOrRef = GenericElementOrRef<Element, Field>
-        return element.hashCode()
-    }
-}

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

@@ -49,7 +49,7 @@ fun field(name: String, type: TypeRef, argument: String? = null, nullable: Boole
     }
 }
 
-fun field(name: String, element: AbstractElement, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(name: String, element: ElementOrRef, nullable: Boolean = false, withReplace: Boolean = false): Field {
     return FirField(name, element, nullable, withReplace)
 }
 
@@ -63,8 +63,8 @@ fun fieldList(name: String, type: TypeRef, withReplace: Boolean = false, useMuta
     return FieldList(name, type, withReplace, useMutableOrEmpty)
 }
 
-fun fieldList(element: AbstractElement, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false): Field {
+fun fieldList(elementOrRef: ElementOrRef, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false): Field {
-    return FieldList(element.name.replaceFirstChar(Char::lowercaseChar) + "s", element, withReplace, useMutableOrEmpty)
+    return FieldList(elementOrRef.element.name.replaceFirstChar(Char::lowercaseChar) + "s", elementOrRef, withReplace, useMutableOrEmpty)
 }
 
 // ----------- Field set -----------

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.fir.tree.generator.printer.generics
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 
 sealed class Field : AbstractField() {
     open var withReplace: Boolean = false
@@ -15,6 +16,12 @@ sealed class Field : AbstractField() {
     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 defaultValueInImplementation: String? get() = null
@@ -40,6 +47,7 @@ sealed class Field : AbstractField() {
             copy.arguments.addAll(arguments)
             copy.needsSeparateTransform = needsSeparateTransform
             copy.needTransformInOtherChildren = needTransformInOtherChildren
+            copy.useInBaseTransformerDetection = useInBaseTransformerDetection
             copy.isMutable = isMutable
             copy.overridenTypes += overridenTypes
             copy.arbitraryImportables += arbitraryImportables
@@ -186,14 +194,12 @@ class SimpleField(
 
 class FirField(
     override val name: String,
-    val element: AbstractElement,
+    val element: ElementOrRef,
     override val nullable: Boolean,
     override var withReplace: Boolean,
 ) : Field() {
     init {
-        if (element is ElementWithArguments) {
+        arguments += element.args.values
-            arguments += element.typeArguments.map { NamedTypeParameterRef(it.name) }
-        }
     }
 
     override val type: String get() = element.type
@@ -254,5 +260,5 @@ class FieldList(
         )
     }
 
-    override val isFirType: Boolean = baseType is AbstractElement
+    override val isFirType: Boolean = baseType is GenericElementOrRef<*, *> && baseType.element is Element
 }

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

@@ -6,27 +6,34 @@
 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.AbstractElement
+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 = mutableMapOf<AbstractElement, Boolean>().withDefault { false }
+    val usedAsFieldType = hashSetOf<GenericElementOrRef<*, *>>()
     for (element in builder.elements) {
         for (field in element.allFirFields) {
+            if (!field.useInBaseTransformerDetection) continue
             val fieldElement = when (field) {
                 is FirField -> field.element
-                is FieldList -> field.baseType as AbstractElement
+                is FieldList -> field.baseType as GenericElementOrRef<*, *>
-                else -> throw IllegalArgumentException()
+                else -> error("Invalid field type: $field")
             }
             if (fieldElement == AbstractFirTreeBuilder.baseFirElement) continue
-            usedAsFieldType[fieldElement] = true
+            usedAsFieldType.add(fieldElement)
         }
     }
 
     for (element in builder.elements) {
         element.traverseParents {
-            if (usedAsFieldType.getValue(it)) {
+            if (it in usedAsFieldType) {
                 element.baseTransformerType = it
                 return@traverseParents
             }

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

@@ -79,6 +79,8 @@ class Element(
     }
 }
 
+// FIXME: Replace this class with `typealias ElementRef = org.jetbrains.kotlin.generators.tree.ElementRef<Element, Field>`
+// FIXME: as soon as Element implements the org.jetbrains.kotlin.generators.tree.AbstractElement interface.
 data class ElementRef(
     val element: Element,
     override val args: Map<NamedTypeParameterRef, TypeRef> = emptyMap(),

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

@@ -10,20 +10,17 @@ import org.jetbrains.kotlin.generators.tree.printer.generics
 /**
  * A common interface representing a FIR or IR tree element.
  */
-interface AbstractElement<Element : AbstractElement<Element, Field>, Field : AbstractField> : FieldContainer, ImplementationKindOwner,
+interface AbstractElement<Element, Field> : ElementOrRef<Element, Field>, FieldContainer, ImplementationKindOwner
-    TypeRef /* TODO: Replace with ElementOrRef */ {
+        where Element : AbstractElement<Element, Field>,
+              Field : AbstractField {
 
     val name: String
 
     val fields: Set<Field>
 
-    val parents: List<Element>
-
     val params: List<TypeVariable>
 
-    // TODO: Remove this property as soon as we replace org.jetbrains.kotlin.fir.tree.generator.model.ElementWithArguments with
+    val parents: List<Element>
-    // a generic ElementRef class
-    val typeArguments: List<TypeArgument>
 
     val parentsArguments: Map<Element, Map<TypeRef, TypeRef>>
 

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

@@ -1,28 +0,0 @@
-/*
- * 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
-
-sealed class TypeArgument(val name: String) {
-    abstract val upperBounds: List<TypeRef>
-}
-
-class SimpleTypeArgument(name: String, val upperBound: TypeRef?) : TypeArgument(name) {
-    override val upperBounds: List<TypeRef> = listOfNotNull(upperBound)
-
-    override fun toString(): String {
-        var result = name
-        if (upperBound != null) {
-            result += " : ${upperBound.typeWithArguments}"
-        }
-        return result
-    }
-}
-
-class TypeArgumentWithMultipleUpperBounds(name: String, override val upperBounds: List<TypeRef>) : TypeArgument(name) {
-    override fun toString(): String {
-        return name
-    }
-}

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

@@ -78,6 +78,41 @@ class ClassRef<P : TypeParameterRef> private constructor(
     override fun toString() = canonicalName
 }
 
+interface ElementOrRef<Element, Field> : ParametrizedTypeRef<ElementOrRef<Element, Field>, NamedTypeParameterRef>, ClassOrElementRef
+        where Element : AbstractElement<Element, Field>,
+              Field : AbstractField {
+    val element: Element
+}
+
+data class ElementRef<Element : AbstractElement<Element, Field>, Field : AbstractField>(
+    override val element: Element,
+    override val args: Map<NamedTypeParameterRef, TypeRef> = emptyMap(),
+    override val nullable: Boolean = false,
+) : ElementOrRef<Element, Field> {
+    override fun copy(args: Map<NamedTypeParameterRef, TypeRef>) = ElementRef(element, args, nullable)
+    override fun copy(nullable: Boolean) = ElementRef(element, args, nullable)
+
+    override val type: String
+        get() = element.type
+    override val packageName: String?
+        get() = element.packageName
+
+    override fun getTypeWithArguments(notNull: Boolean): String {
+        return element.type + generics
+    }
+
+    override fun toString() = buildString {
+        append(element.name)
+        append("<")
+        append(args)
+        append(">")
+        if (nullable) {
+            append("?")
+        }
+    }
+}
+
+
 sealed interface TypeParameterRef : TypeRef
 
 data class PositionTypeParameterRef(