From d904ce0939fd609339e54a257d9c64afdffdf1ce Mon Sep 17 00:00:00 2001
From: Wojciech Litewka <wojciech.litewka@jetbrains.com>
Date: Wed, 21 Feb 2024 17:33:36 +0100
Subject: [PATCH] [tree generator] Merge default and custom implementations

---
 .../kotlin/generators/tree/AbstractElement.kt      | 14 +-------------
 .../generators/tree/AbstractImplementation.kt      |  9 +--------
 .../tree/config/AbstractBuilderConfigurator.kt     | 12 ++++++------
 .../config/AbstractImplementationConfigurator.kt   |  9 +++------
 .../kotlin/generators/tree/printer/common.kt       |  2 +-
 5 files changed, 12 insertions(+), 34 deletions(-)

diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractElement.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractElement.kt
index ecd230b4afd..a2821da5c18 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractElement.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractElement.kt
@@ -154,22 +154,10 @@ abstract class AbstractElement<Element, Field, Implementation>(
     val transformerClass: Element
         get() = transformerReturnType ?: baseTransformerType ?: element
 
-    var defaultImplementation: Implementation? = null
-
-    val customImplementations = mutableListOf<Implementation>()
+    val implementations = mutableListOf<Implementation>()
 
     var doesNotNeedImplementation: Boolean = false
 
-    val allImplementations: List<Implementation> by lazy {
-        if (doesNotNeedImplementation) {
-            emptyList()
-        } else {
-            val implementations = customImplementations.toMutableList()
-            defaultImplementation?.let { implementations += it }
-            implementations
-        }
-    }
-
     /**
      * Types/functions that you want to additionally import in the file with the element class.
      *
diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractImplementation.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractImplementation.kt
index 7ca12e9d77d..c813c58d3b6 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractImplementation.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractImplementation.kt
@@ -17,9 +17,6 @@ abstract class AbstractImplementation<Implementation, Element, Field>(
               Element : AbstractElement<Element, *, Implementation>,
               Field : AbstractField<*> {
 
-    private val isDefault: Boolean
-        get() = name == null
-
     override val allParents: List<ImplementationKindOwner>
         get() = listOf(element)
 
@@ -53,11 +50,7 @@ abstract class AbstractImplementation<Implementation, Element, Field>(
 
     init {
         @Suppress("UNCHECKED_CAST")
-        if (isDefault) {
-            element.defaultImplementation = this as Implementation
-        } else {
-            element.customImplementations += this as Implementation
-        }
+        element.implementations += this as Implementation
     }
 
     override val hasAcceptChildrenMethod: Boolean
diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractBuilderConfigurator.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractBuilderConfigurator.kt
index 02874dd3f68..3e7061087f9 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractBuilderConfigurator.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractBuilderConfigurator.kt
@@ -85,10 +85,10 @@ abstract class AbstractBuilderConfigurator<Element, Implementation, BuilderField
 
     private fun Element.extractImplementation(type: String?): Implementation {
         return if (type == null) {
-            allImplementations.singleOrNull { it.kind?.hasLeafBuilder == true } ?: this@AbstractBuilderConfigurator.run {
+            implementations.singleOrNull { it.kind?.hasLeafBuilder == true } ?: this@AbstractBuilderConfigurator.run {
                 val message = buildString {
                     appendLine("${this@extractImplementation} has multiple implementations:")
-                    for (implementation in allImplementations) {
+                    for (implementation in implementations) {
                         appendLine("  - ${implementation.typeName}")
                     }
                     appendLine("Please specify implementation is needed")
@@ -96,10 +96,10 @@ abstract class AbstractBuilderConfigurator<Element, Implementation, BuilderField
                 throw IllegalArgumentException(message)
             }
         } else {
-            allImplementations.firstOrNull { it.typeName == type } ?: this@AbstractBuilderConfigurator.run {
+            implementations.firstOrNull { it.typeName == type } ?: this@AbstractBuilderConfigurator.run {
                 val message = buildString {
                     appendLine("${this@extractImplementation} has not implementation $type. Existing implementations:")
-                    for (implementation in allImplementations) {
+                    for (implementation in implementations) {
                         appendLine("  - ${implementation.typeName}")
                     }
                     appendLine("Please specify implementation is needed")
@@ -118,7 +118,7 @@ abstract class AbstractBuilderConfigurator<Element, Implementation, BuilderField
         implementationPredicate: (Implementation) -> Boolean = { true }
     ): Collection<Implementation> {
         return elements
-            .flatMap { it.allImplementations }
+            .flatMap { it.implementations }
             .mapNotNullTo(mutableSetOf()) { implementation ->
                 if (!implementationPredicate(implementation)) return@mapNotNullTo null
                 if (implementation.element == element) return@mapNotNullTo null
@@ -128,7 +128,7 @@ abstract class AbstractBuilderConfigurator<Element, Implementation, BuilderField
     }
 
     private val allLeafBuilders: List<LeafBuilder<BuilderField, Element, Implementation>>
-        get() = elements.flatMap { it.allImplementations }.mapNotNull { it.builder }
+        get() = elements.flatMap { it.implementations }.mapNotNull { it.builder }
 
     /**
      * Allows to batch-apply [config] to certain fields in _all_ the builders that satisfy the given
diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractImplementationConfigurator.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractImplementationConfigurator.kt
index 616f2af1ce4..8517ac41569 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractImplementationConfigurator.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/config/AbstractImplementationConfigurator.kt
@@ -60,11 +60,8 @@ abstract class AbstractImplementationConfigurator<Implementation, Element, Imple
      * @return The configured implementation.
      */
     protected fun impl(element: Element, name: String? = null, config: ImplementationContext.() -> Unit = {}): Implementation {
-        val implementation = if (name == null) {
-            element.defaultImplementation
-        } else {
-            element.customImplementations.firstOrNull { it.name == name }
-        } ?: createImplementation(element, name)
+        val implementation = element.implementations.firstOrNull { it.name == name }
+            ?: createImplementation(element, name)
         val context = ImplementationContext(implementation)
         context.apply(config)
         elementsWithImpl += element
@@ -103,7 +100,7 @@ abstract class AbstractImplementationConfigurator<Implementation, Element, Imple
         config: ImplementationContext.(field: String) -> Unit,
     ) {
         for (element in elementsWithImpl) {
-            for (implementation in element.allImplementations) {
+            for (implementation in element.implementations) {
                 if (!implementationPredicate(implementation)) continue
                 if (!implementation.allFields.any { it.name == field }) continue
                 if (!fieldPredicate(implementation.getField(field))) continue
diff --git a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/common.kt b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/common.kt
index 76ad17aa15f..0d1d5c9ee0d 100644
--- a/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/common.kt
+++ b/generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/common.kt
@@ -101,7 +101,7 @@ fun <Element, Implementation, ElementField, ImplementationField> generateTree(
         detectBaseTransformerTypes(model)
     }
     implementationConfigurator?.configureImplementations(model)
-    val implementations = model.elements.flatMap { it.allImplementations }
+    val implementations = model.elements.flatMap { it.implementations }
     InterfaceAndAbstractClassConfigurator((model.elements + implementations))
         .configureInterfacesAndAbstractClasses()
     addPureAbstractElement(model.elements, pureAbstractElement)