[FIR/IR generator] Factor out allFields to AbstractElement

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

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

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

@@ -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>
+    override val overriddenTypes: MutableSet<TypeRefWithNullability>
-        get() = origin.overridenTypes
+        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,

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

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

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

@@ -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)")

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

@@ -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)

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

@@ -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)

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

@@ -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,

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

@@ -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,

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

@@ -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) {

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

@@ -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)

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

@@ -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 {

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

@@ -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.

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

@@ -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) {
+                    for (field in filterFields(element)) {
-                        if (
+                        if (field.isParameter) continue
-                            !field.withGetter && field.defaultValueInImplementation == null && field.isFinal && field.fromParent ||
+                        if (!field.withGetter && field.defaultValueInImplementation == null && field.isFinal && field.fromParent) {
-                            field.isParameter
-                        ) {
                             continue
                         }
                         if (separateFieldsWithBlankLine) println()

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

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

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

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

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

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

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

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