From 075068155bb11d349f993305bea7bfebf9eada69 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 30 Jan 2020 13:26:37 +0300
Subject: [PATCH] [FIR] Delay configuration of nodes in tree generator

---
 .../kotlin/fir/tree/generator/NodeConfigurator.kt     |  4 ++--
 .../generator/context/AbstractFieldConfigurator.kt    | 11 ++++++++---
 .../tree/generator/context/AbstractFirTreeBuilder.kt  |  8 ++++++++
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt
index ba5b16f5599..bfb870590ed 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt
+++ b/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() {
-    fun configureFields() = with(FirTreeBuilder) {
+object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuilder) {
+    fun configureFields() = configure {
         AbstractFirTreeBuilder.baseFirElement.configure {
             +field("source", sourceElementType, nullable = true)
         }
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFieldConfigurator.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFieldConfigurator.kt
index dcb80bfd755..697b6a0cf1d 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFieldConfigurator.kt
+++ b/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) {
-        ConfigureContext(this).block()
+    fun Element.configure(block: ConfigureContext.() -> Unit) {
+        builder.configurations[this] = { ConfigureContext(this).block() }
+    }
+
+    fun configure(init: T.() -> Unit) {
+        builder.init()
+        builder.applyConfigurations()
     }
 }
\ No newline at end of file
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFirTreeBuilder.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFirTreeBuilder.kt
index b8c1bd4261e..715324bd7de 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractFirTreeBuilder.kt
+++ b/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)