[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.
2023-09-14 15:27:36 +02:00
parent 86c8f3cc35
commit 04ac4b71a3
10 changed files with 98 additions and 97 deletions
@@ -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(
@@ -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)
-}
@@ -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> {
-    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 {
+class Element(override val name: String, kind: Kind) : AbstractElement<Element, Field> {
    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
-    override val customImplementations = mutableListOf<Implementation>()
+    var defaultImplementation: Implementation? = null
+    val customImplementations = mutableListOf<Implementation>()

    override val params = mutableListOf<TypeVariable>()

-    override val typeArguments: List<TypeArgument>
-        get() = emptyList()
-
-    override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<TypeRef, TypeRef>>()
+    override val parentsArguments = mutableMapOf<Element, 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
-    override val transformerType: Element get() = baseTransformerType ?: this
+    var baseTransformerType: Element? = null
+    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()
-    override val allImplementations: List<Implementation> by lazy {
+    val useNullableForReplace: MutableSet<Field> = mutableSetOf()
+    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 {
-    override fun equals(other: Any?): Boolean {
-        return element.equals(other)
-    }
+typealias ElementRef = GenericElementRef<Element, Field>

-    override fun hashCode(): Int {
-        return element.hashCode()
-    }
-}
+typealias ElementOrRef = GenericElementOrRef<Element, Field>
@@ -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 {
-    return FieldList(element.name.replaceFirstChar(Char::lowercaseChar) + "s", element, withReplace, useMutableOrEmpty)
+fun fieldList(elementOrRef: ElementOrRef, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false): Field {
+    return FieldList(elementOrRef.element.name.replaceFirstChar(Char::lowercaseChar) + "s", elementOrRef, withReplace, useMutableOrEmpty)
 }

 // ----------- Field set -----------
@@ -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.typeArguments.map { NamedTypeParameterRef(it.name) }
-        }
+        arguments += element.args.values
    }

    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
 }
@@ -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
-                else -> throw IllegalArgumentException()
+                is FieldList -> field.baseType as GenericElementOrRef<*, *>
+                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
            }
@@ -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(),
@@ -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,
-    TypeRef /* TODO: Replace with ElementOrRef */ {
+interface AbstractElement<Element, Field> : ElementOrRef<Element, Field>, FieldContainer, ImplementationKindOwner
+        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
-    // a generic ElementRef class
-    val typeArguments: List<TypeArgument>
+    val parents: List<Element>

    val parentsArguments: Map<Element, Map<TypeRef, TypeRef>>

@@ -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
-    }
-}
@@ -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(