[FIR] Delay configuration of nodes in tree generator

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

@@ -32,8 +32,8 @@ import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 
-object NodeConfigurator : AbstractFieldConfigurator() {
+object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuilder) {
-    fun configureFields() = with(FirTreeBuilder) {
+    fun configureFields() = configure {
         AbstractFirTreeBuilder.baseFirElement.configure {
             +field("source", sourceElementType, nullable = true)
         }

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

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.context
 
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 
-abstract class AbstractFieldConfigurator {
+abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val builder: T) {
     inner class ConfigureContext(val element: Element) {
         operator fun FieldSet.unaryPlus() {
             element.fields.addAll(this.map { it.copy() })
@@ -78,7 +78,12 @@ abstract class AbstractFieldConfigurator {
         }
     }
 
-    inline fun Element.configure(block: ConfigureContext.() -> Unit) {
+    fun Element.configure(block: ConfigureContext.() -> Unit) {
-        ConfigureContext(this).block()
+        builder.configurations[this] = { ConfigureContext(this).block() }
+    }
+
+    fun configure(init: T.() -> Unit) {
+        builder.init()
+        builder.applyConfigurations()
     }
 }

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

@@ -32,6 +32,14 @@ abstract class AbstractFirTreeBuilder {
             it.parents.addAll(dependencies)
             elements += it
         }
+
+    val configurations: MutableMap<Element, () -> Unit> = mutableMapOf()
+
+    fun applyConfigurations() {
+        for (element in elements) {
+            configurations[element]?.invoke()
+        }
+    }
 }
 
 fun generatedType(type: String): Type = generatedType("", type)