[FIR/IR generator] Factor out allFields to AbstractElement

2023-11-14 19:45:41 +01:00
parent 38ee477ab4
commit ad4e97154f
17 changed files with 153 additions and 178 deletions
@@ -7,7 +7,6 @@ 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.printer.safeDecapitalizedName
-import org.jetbrains.kotlin.fir.tree.generator.util.set
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef
@@ -80,8 +79,6 @@ class Element(name: String, override val propertyName: String, kind: Kind) : Abs
    var doesNotNeedImplementation: Boolean = false

    val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || elementParents.any { it.element.needTransformOtherChildren }
-    val overridenFields: MutableMap<Field, MutableMap<Field, Boolean>> = mutableMapOf()
-    val useNullableForReplace: MutableSet<Field> = mutableSetOf()
    val allImplementations: List<Implementation> by lazy {
        if (doesNotNeedImplementation) {
            emptyList()
@@ -92,66 +89,6 @@ class Element(name: String, override val propertyName: String, kind: Kind) : Abs
        }
    }

-    override val allFields: List<Field> by lazy {
-        val result = LinkedHashSet<Field>()
-        result.addAll(fields.toList().asReversed())
-        result.forEach { overridenFields[it, it] = false }
-        for (parentField in parentFields.asReversed()) {
-            val overrides = !result.add(parentField)
-            if (overrides) {
-                val existingField = result.first { it == parentField }
-                existingField.fromParent = true
-                existingField.needsSeparateTransform = existingField.needsSeparateTransform || parentField.needsSeparateTransform
-                existingField.needTransformInOtherChildren = existingField.needTransformInOtherChildren || parentField.needTransformInOtherChildren
-                existingField.withReplace = parentField.withReplace || existingField.withReplace
-                existingField.parentHasSeparateTransform = parentField.needsSeparateTransform
-                if (parentField.typeRef.copy(nullable = false) != existingField.typeRef.copy(nullable = false) && parentField.withReplace) {
-                    existingField.overridenTypes += parentField.typeRef
-                    overridenFields[existingField, parentField] = false
-                } else {
-                    overridenFields[existingField, parentField] = true
-                    if (parentField.nullable != existingField.nullable) {
-                        existingField.useNullableForReplace = true
-                    }
-                }
-            } else {
-                overridenFields[parentField, parentField] = true
-            }
-        }
-        result.toList().asReversed()
-    }
-
-    val parentFields: List<Field> by lazy {
-        val result = LinkedHashMap<String, Field>()
-        elementParents.forEach { parentRef ->
-            val parent = parentRef.element
-            val fields = parent.allFields.map { field ->
-                val copy = (field as? SimpleField)?.let { simpleField ->
-                    simpleField.replaceType(simpleField.typeRef.substitute(parentRef.args) as TypeRefWithNullability)
-                } ?: field.copy()
-                copy.apply {
-                    fromParent = true
-                }
-            }
-            fields.forEach {
-                result.merge(it.name, it) { previousField, thisField ->
-                    val resultField = previousField.copy()
-                    if (thisField.withReplace) {
-                        resultField.withReplace = true
-                    }
-                    if (thisField.useNullableForReplace) {
-                        resultField.useNullableForReplace = true
-                    }
-                    if (thisField.isMutable) {
-                        resultField.isMutable = true
-                    }
-                    resultField
-                }
-            }
-        }
-        result.values.toList()
-    }
-
    override fun toString(): String {
        return with(ImportCollector("")) { render() }
    }
@@ -22,7 +22,6 @@ sealed class Field : AbstractField<Field>() {
    open var isMutableInInterface: Boolean = false
    open val fromDelegate: Boolean get() = false

-    open val overridenTypes: MutableSet<TypeRefWithNullability> = mutableSetOf()
    open var useNullableForReplace: Boolean = false
    open var notNull: Boolean = false

@@ -38,7 +37,9 @@ sealed class Field : AbstractField<Field>() {

    abstract override var isMutable: Boolean

-    fun copy(): Field = internalCopy().also {
+    override fun replaceType(newType: TypeRefWithNullability): Field = copy()
+
+    override fun copy(): Field = internalCopy().also {
        updateFieldsInCopy(it)
    }

@@ -48,7 +49,7 @@ sealed class Field : AbstractField<Field>() {
            copy.needTransformInOtherChildren = needTransformInOtherChildren
            copy.useInBaseTransformerDetection = useInBaseTransformerDetection
            copy.isMutable = isMutable
-            copy.overridenTypes += overridenTypes
+            copy.overriddenTypes += overriddenTypes
            copy.arbitraryImportables += arbitraryImportables
            copy.useNullableForReplace = useNullableForReplace
            copy.customInitializationCall = customInitializationCall
@@ -61,6 +62,15 @@ sealed class Field : AbstractField<Field>() {

    protected abstract fun internalCopy(): Field

+    override fun updatePropertiesFromOverriddenField(parentField: Field, haveSameClass: Boolean) {
+        needsSeparateTransform = needsSeparateTransform || parentField.needsSeparateTransform
+        needTransformInOtherChildren = needTransformInOtherChildren || parentField.needTransformInOtherChildren
+        withReplace = withReplace || parentField.withReplace
+        parentHasSeparateTransform = parentField.needsSeparateTransform
+        if (parentField.nullable != nullable && haveSameClass) {
+            useNullableForReplace = true
+        }
+    }
 }

 // ----------- Field with default -----------
@@ -112,8 +122,8 @@ class FieldWithDefault(val origin: Field) : Field() {
    override var withGetter: Boolean = false
    override var customSetter: String? = null
    override var fromDelegate: Boolean = false
-    override val overridenTypes: MutableSet<TypeRefWithNullability>
-        get() = origin.overridenTypes
+    override val overriddenTypes: MutableSet<TypeRefWithNullability>
+        get() = origin.overriddenTypes

    override val arbitraryImportables: MutableList<Importable>
        get() = origin.arbitraryImportables
@@ -158,7 +168,7 @@ class SimpleField(
        }
    }

-    fun replaceType(newType: TypeRefWithNullability) = SimpleField(
+    override fun replaceType(newType: TypeRefWithNullability) = SimpleField(
        name = name,
        typeRef = newType,
        withReplace = withReplace,
@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.tree.generator.*
 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.util.get
+import org.jetbrains.kotlin.generators.tree.get
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.printer.*
 import org.jetbrains.kotlin.utils.SmartPrinter
@@ -43,9 +43,9 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
            }

            allFields.filter { it.withReplace }.forEach {
-                val override = overridenFields[it, it] && !(it.name == "source" && element == FirTreeBuilder.qualifiedAccessExpression)
+                val override = overriddenFieldsHaveSameClass[it, it] && !(it.name == "source" && element == FirTreeBuilder.qualifiedAccessExpression)
                it.replaceDeclaration(override, forceNullable = it.useNullableForReplace)
-                for (overriddenType in it.overridenTypes) {
+                for (overriddenType in it.overriddenTypes) {
                    it.replaceDeclaration(true, overriddenType)
                }
            }
@@ -380,7 +380,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                    }
                }

-                for (overridenType in field.overridenTypes) {
+                for (overridenType in field.overriddenTypes) {
                    generateReplace(field, overridenType) {
                        println("require($newValue is ${field.typeRef.render()})")
                        println("replace$capitalizedFieldName($newValue)")
@@ -9,16 +9,6 @@ import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil
 import org.jetbrains.kotlin.utils.SmartPrinter
 import java.io.File

-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)
-    map[k2] = value
-}
-
-operator fun <K, V, U> Map<K, Map<V, U>>.get(k1: K, k2: V): U {
-    return getValue(k1).getValue(k2)
-}
-
 inline fun File.writeToFileUsingSmartPrinterIfFileContentChanged(block: SmartPrinter.() -> Unit) {
    val newText = buildString { SmartPrinter(this).block() }
    GeneratorsFileUtil.writeFileIfContentChanged(this, newText, logNotChanged = false)
@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil.collectPreviously
 import org.jetbrains.kotlin.generators.util.GeneratorsFileUtil.removeExtraFilesFromPreviousGeneration
 import org.jetbrains.kotlin.ir.generator.model.addWalkableChildren
 import org.jetbrains.kotlin.ir.generator.model.markLeaves
-import org.jetbrains.kotlin.ir.generator.model.processFieldOverrides
 import org.jetbrains.kotlin.ir.generator.print.*
 import java.io.File

@@ -29,7 +28,6 @@ fun main(args: Array<String>) {
    InterfaceAndAbstractClassConfigurator(model.elements).configureInterfacesAndAbstractClasses()
    addPureAbstractElement(model.elements, elementBaseType)
    markLeaves(model.elements)
-    processFieldOverrides(model.elements)
    addWalkableChildren(model.elements)

    val previouslyGeneratedFiles = collectPreviouslyGeneratedFiles(generationPath)
@@ -72,19 +72,19 @@ abstract class AbstractTreeBuilder {

    protected fun field(
        name: String,
-        type: TypeRefWithNullability?,
+        type: TypeRefWithNullability,
        nullable: Boolean = false,
        mutable: Boolean = true,
        isChild: Boolean = false,
        initializer: SingleField.() -> Unit = {}
    ): SingleField {
        checkChildType(isChild, type, name)
-        return SingleField(name, type?.copy(nullable) ?: InferredOverriddenType, mutable, isChild).apply(initializer)
+        return SingleField(name, type.copy(nullable), mutable, isChild).apply(initializer)
    }

    protected fun listField(
        name: String,
-        elementType: TypeRef?,
+        elementType: TypeRef,
        nullable: Boolean = false,
        mutability: ListField.Mutability,
        isChild: Boolean = false,
@@ -98,7 +98,7 @@ abstract class AbstractTreeBuilder {
        }
        return ListField(
            name = name,
-            baseType = elementType ?: InferredOverriddenType,
+            baseType = elementType,
            listType = listType,
            isNullable = nullable,
            mutable = mutability == ListField.Mutability.Var,
@@ -42,10 +42,6 @@ class Element(
    var generateIrFactoryMethod = category == Category.Declaration
    val fieldsToSkipInIrFactoryMethod = hashSetOf<String>()

-
-    override val allFields: List<Field>
-        get() = fields.toList()
-
    override val element: Element
        get() = this

@@ -124,29 +120,6 @@ class Element(

    override fun toString() = name

-    fun elementParentsRecursively(): List<ElementRef> {
-        val linkedSet = buildSet {
-            fun recurse(element: Element) {
-                addAll(element.elementParents)
-                element.elementParents.forEach { recurse(it.element) }
-            }
-            recurse(this@Element)
-        }
-        return topologicalSort(linkedSet) {
-            element.elementParents
-        }
-    }
-
-    fun allFieldsRecursively(): List<Field> {
-        val parentFields = elementParentsRecursively()
-            .reversed()
-            .flatMap { it.element.fields }
-        return (parentFields + fields)
-            .asReversed()
-            .distinctBy { it.name }
-            .asReversed()
-    }
-
    operator fun TypeVariable.unaryPlus() = apply {
        params.add(this)
    }
@@ -198,6 +171,24 @@ sealed class Field(

    override val isParameter: Boolean
        get() = false
+
+    override fun copy() = internalCopy().also(::updateFieldsInCopy)
+
+    protected fun updateFieldsInCopy(copy: Field) {
+        copy.useInBaseTransformerDetection = useInBaseTransformerDetection
+        copy.isMutable = isMutable
+        copy.optInAnnotation = optInAnnotation
+        copy.baseDefaultValue = baseDefaultValue
+        copy.baseGetter = baseGetter
+        copy.useInIrFactoryStrategy = useInIrFactoryStrategy
+        copy.fromParent = fromParent
+        copy.customSetter = customSetter
+        copy.visibility = visibility
+        copy.deprecation = deprecation
+        copy.kDoc = kDoc
+    }
+
+    protected abstract fun internalCopy(): Field
 }

 class SingleField(
@@ -205,7 +196,13 @@ class SingleField(
    override var typeRef: TypeRefWithNullability,
    mutable: Boolean,
    isChild: Boolean,
-) : Field(name, mutable, isChild)
+) : Field(name, mutable, isChild) {
+
+    override fun replaceType(newType: TypeRefWithNullability) =
+        SingleField(name, newType, isMutable, isChild).also(::updateFieldsInCopy)
+
+    override fun internalCopy() = SingleField(name, typeRef, isMutable, isChild)
+}

 class ListField(
    name: String,
@@ -219,6 +216,10 @@ class ListField(
    override val typeRef: ClassRef<PositionTypeParameterRef>
        get() = listType.withArgs(baseType).copy(isNullable)

+    override fun replaceType(newType: TypeRefWithNullability) = copy()
+
+    override fun internalCopy() = ListField(name, baseType, isNullable, listType, isMutable, isChild)
+
    enum class Mutability {
        Var,
        MutableList,
@@ -7,21 +7,6 @@ package org.jetbrains.kotlin.ir.generator.model

 import org.jetbrains.kotlin.generators.tree.*

-internal object InferredOverriddenType : TypeRefWithNullability {
-
-    context(ImportCollector)
-    override fun renderTo(appendable: Appendable) {
-        renderingIsNotSupported()
-    }
-
-    override fun substitute(map: TypeParameterSubstitutionMap) = this
-
-    override val nullable: Boolean
-        get() = false
-
-    override fun copy(nullable: Boolean) = this
-}
-
 typealias Model = org.jetbrains.kotlin.generators.tree.Model<Element>

 internal fun markLeaves(elements: List<Element>) {
@@ -40,39 +25,6 @@ internal fun markLeaves(elements: List<Element>) {
    }
 }

-internal fun processFieldOverrides(elements: List<Element>) {
-    for (element in iterateElementsParentFirst(elements)) {
-        for (field in element.fields) {
-            fun visitParents(visited: Element) {
-                for (parent in visited.elementParents) {
-                    val overriddenField = parent.element.fields.singleOrNull { it.name == field.name }
-                    if (overriddenField != null) {
-                        field.fromParent = true
-                        field.optInAnnotation = field.optInAnnotation ?: overriddenField.optInAnnotation
-
-                        fun transformInferredType(type: TypeRef, overriddenType: TypeRef) =
-                            type.takeUnless { it is InferredOverriddenType } ?: overriddenType
-                        when (field) {
-                            is SingleField -> {
-                                field.typeRef =
-                                    transformInferredType(field.typeRef, (overriddenField as SingleField).typeRef) as TypeRefWithNullability
-                            }
-                            is ListField -> {
-                                field.baseType = transformInferredType(field.baseType, (overriddenField as ListField).baseType)
-                            }
-                        }
-
-                        break
-                    }
-
-                    visitParents(parent.element)
-                }
-            }
-
-            visitParents(element)
-        }
-    }
-}

 internal fun addWalkableChildren(elements: List<Element>) {
    for (element in elements) {
@@ -11,11 +11,12 @@ 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.ir.generator.model.Element
+import org.jetbrains.kotlin.ir.generator.model.Field
 import org.jetbrains.kotlin.ir.generator.model.ListField
 import org.jetbrains.kotlin.ir.generator.model.Model
+import org.jetbrains.kotlin.ir.generator.model.SingleField
 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

@@ -34,6 +35,16 @@ private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Ele
    override val separateFieldsWithBlankLine: Boolean
        get() = true

+    // In IR classes we only print fields that are either declared in this element, or refine the type of a parent field
+    // and thus need an override.
+    override fun filterFields(element: Element): Collection<Field> = element.fields.map { field ->
+        field.copy().apply {
+            if (field in element.parentFields) {
+                fromParent = true
+            }
+        }
+    }
+
    context(ImportCollector)
    override fun SmartPrinter.printAdditionalMethods(element: Element) {
        element.generationCallback?.invoke(this@ImportCollector, this)
@@ -255,7 +255,7 @@ private class FactoryMethod(val element: Element) {

    val name = "create" + element.name.capitalizeAsciiOnly()

-    val parameters = (element.allFieldsRecursively() + element.additionalIrFactoryMethodParameters)
+    val parameters = (element.allFields + element.additionalIrFactoryMethodParameters)
        .filterNot { it.name in element.fieldsToSkipInIrFactoryMethod }
        .mapNotNull { field ->
            (field.useInIrFactoryStrategy as? Field.UseFieldAsParameterInIrFactoryStrategy.Yes)?.let {
@@ -78,7 +78,52 @@ abstract class AbstractElement<Element, Field>(
        }
    }

-    abstract override val allFields: List<Field>
+    override val allFields: List<Field> by lazy {
+        val result = LinkedHashSet<Field>()
+        result.addAll(fields.toList().asReversed())
+        result.forEach { overriddenFieldsHaveSameClass[it, it] = false }
+        for (parentField in parentFields.asReversed()) {
+            val overrides = !result.add(parentField)
+            if (overrides) {
+                val existingField = result.first { it == parentField }
+                existingField.fromParent = true
+                val haveSameClass = parentField.typeRef.copy(nullable = false) == existingField.typeRef.copy(nullable = false)
+                if (!haveSameClass) {
+                    existingField.overriddenTypes += parentField.typeRef
+                }
+                overriddenFieldsHaveSameClass[existingField, parentField] = haveSameClass
+                existingField.updatePropertiesFromOverriddenField(parentField, haveSameClass)
+            } else {
+                overriddenFieldsHaveSameClass[parentField, parentField] = true
+            }
+        }
+        result.toList().asReversed()
+    }
+
+    val overriddenFieldsHaveSameClass: MutableMap<Field, MutableMap<Field, Boolean>> = mutableMapOf()
+
+    val parentFields: List<Field> by lazy {
+        val result = LinkedHashMap<String, Field>()
+        elementParents.forEach { parentRef ->
+            val parent = parentRef.element
+            val fields = parent.allFields.map { field ->
+                field.replaceType(field.typeRef.substitute(parentRef.args) as TypeRefWithNullability)
+                    .apply {
+                        fromParent = true
+                    }
+            }
+            fields.forEach {
+                result.merge(it.name, it) { previousField, thisField ->
+                    val resultField = previousField.copy()
+                    if (thisField.isMutable) {
+                        resultField.isMutable = true
+                    }
+                    resultField
+                }
+            }
+        }
+        result.values.toList()
+    }

    /**
     * A custom return type of the corresponding transformer method for this element.
@@ -27,6 +27,8 @@ abstract class AbstractElementPrinter<Element : AbstractElement<Element, Field>,
    protected open val separateFieldsWithBlankLine: Boolean
        get() = false

+    protected open fun filterFields(element: Element): Collection<Field> = element.allFields
+
    context(ImportCollector)
    fun printElement(element: Element) {
        printer.run {
@@ -49,11 +51,9 @@ abstract class AbstractElementPrinter<Element : AbstractElement<Element, Field>,
            val body = SmartPrinter(StringBuilder()).apply {
                val fieldPrinter = makeFieldPrinter(this)
                withIndent {
-                    for (field in element.allFields) {
-                        if (
-                            !field.withGetter && field.defaultValueInImplementation == null && field.isFinal && field.fromParent ||
-                            field.isParameter
-                        ) {
+                    for (field in filterFields(element)) {
+                        if (field.isParameter) continue
+                        if (!field.withGetter && field.defaultValueInImplementation == null && field.isFinal && field.fromParent) {
                            continue
                        }
                        if (separateFieldsWithBlankLine) println()
@@ -28,7 +28,7 @@ abstract class AbstractField<Field : AbstractField<Field>> {

    open var optInAnnotation: ClassRef<*>? = null

-    abstract val isMutable: Boolean
+    abstract var isMutable: Boolean
    open val withGetter: Boolean get() = false
    open val customSetter: String? get() = null

@@ -58,6 +58,10 @@ abstract class AbstractField<Field : AbstractField<Field>> {
     */
    var needAcceptAndTransform: Boolean = true

+    open val overriddenTypes: MutableSet<TypeRefWithNullability> = mutableSetOf()
+
+    open fun updatePropertiesFromOverriddenField(parentField: Field, haveSameClass: Boolean) {}
+
    override fun toString(): String {
        return name
    }
@@ -73,4 +77,14 @@ abstract class AbstractField<Field : AbstractField<Field>> {
    override fun hashCode(): Int {
        return name.hashCode()
    }
+
+    /**
+     * Returns a copy of this field with its [typeRef] set to [newType] (if it's possible).
+     */
+    abstract fun replaceType(newType: TypeRefWithNullability): Field
+
+    /**
+     * Returns a copy of this field.
+     */
+    abstract fun copy(): Field
 }
@@ -34,7 +34,9 @@ abstract class AbstractFieldPrinter<Field : AbstractField<Field>>(
        modality: Modality? = null,
    ) {
        printer.run {
-            printKDoc(field.kDoc)
+            if (!field.fromParent) {
+                printKDoc(field.kDoc)
+            }

            field.deprecation?.let {
                println("@Deprecated(")
@@ -107,4 +109,4 @@ abstract class AbstractFieldPrinter<Field : AbstractField<Field>>(
            }
        }
    }
-}
+}
@@ -6,7 +6,12 @@
 package org.jetbrains.kotlin.generators.tree

 interface FieldContainer<out Field : AbstractField<*>> {
+
+    /**
+     * All the fields of this element, including the fields of all its parents.
+     */
    val allFields: List<Field>
+
    operator fun get(fieldName: String): Field?

    val hasAcceptMethod: Boolean
@@ -32,4 +37,4 @@ interface FieldContainer<out Field : AbstractField<*>> {
     */
    val transformableChildren: List<Field>
        get() = walkableChildren.filter { it.isMutable || it is ListField }
-}
+}
@@ -36,4 +36,14 @@ fun <Element : AbstractElement<Element, *>> detectBaseTransformerTypes(model: Mo
            }
        }
    }
+}
+
+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)
+    map[k2] = value
+}
+
+operator fun <K, V, U> Map<K, Map<V, U>>.get(k1: K, k2: V): U {
+    return getValue(k1).getValue(k2)
 }