From 8ffb3297418b9db079ba426cb946bd7bb9a4e6f3 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Fri, 23 Jun 2023 21:56:47 +0900
Subject: [PATCH] [FIR] Add 'buildCodeFragmentCopy()' and 'replaceBlock()' for
 IDE analysis

---
 .../kotlin/fir/declarations/FirCodeFragment.kt |  2 ++
 .../builder/FirCodeFragmentBuilder.kt          | 18 ++++++++++++++++++
 .../declarations/impl/FirCodeFragmentImpl.kt   |  4 ++++
 .../fir/tree/generator/BuilderConfigurator.kt  |  4 ++++
 .../fir/tree/generator/NodeConfigurator.kt     |  2 +-
 5 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirCodeFragment.kt b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirCodeFragment.kt
index 3465196a36c..c1d4681e867 100644
--- a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirCodeFragment.kt
+++ b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirCodeFragment.kt
@@ -36,6 +36,8 @@ abstract class FirCodeFragment : FirDeclaration() {
 
     abstract override fun replaceAnnotations(newAnnotations: List<FirAnnotation>)
 
+    abstract fun replaceBlock(newBlock: FirBlock)
+
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirCodeFragment
 
     abstract fun <D> transformBlock(transformer: FirTransformer<D>, data: D): FirCodeFragment
diff --git a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirCodeFragmentBuilder.kt b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirCodeFragmentBuilder.kt
index 0eadd6674ec..b063e02a4dd 100644
--- a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirCodeFragmentBuilder.kt
+++ b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirCodeFragmentBuilder.kt
@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolveState
 import org.jetbrains.kotlin.fir.declarations.ResolveStateAccess
 import org.jetbrains.kotlin.fir.declarations.asResolveState
 import org.jetbrains.kotlin.fir.declarations.impl.FirCodeFragmentImpl
+import org.jetbrains.kotlin.fir.declarations.resolvePhase
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirBlock
 import org.jetbrains.kotlin.fir.symbols.impl.FirCodeFragmentSymbol
@@ -64,3 +65,20 @@ inline fun buildCodeFragment(init: FirCodeFragmentBuilder.() -> Unit): FirCodeFr
     }
     return FirCodeFragmentBuilder().apply(init).build()
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildCodeFragmentCopy(original: FirCodeFragment, init: FirCodeFragmentBuilder.() -> Unit): FirCodeFragment {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    val copyBuilder = FirCodeFragmentBuilder()
+    copyBuilder.source = original.source
+    copyBuilder.resolvePhase = original.resolvePhase
+    copyBuilder.annotations.addAll(original.annotations)
+    copyBuilder.moduleData = original.moduleData
+    copyBuilder.origin = original.origin
+    copyBuilder.attributes = original.attributes.copy()
+    copyBuilder.symbol = original.symbol
+    copyBuilder.block = original.block
+    return copyBuilder.apply(init).build()
+}
diff --git a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirCodeFragmentImpl.kt b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirCodeFragmentImpl.kt
index bc2c256d9de..89b478c3230 100644
--- a/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirCodeFragmentImpl.kt
+++ b/compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirCodeFragmentImpl.kt
@@ -68,4 +68,8 @@ internal class FirCodeFragmentImpl(
     override fun replaceAnnotations(newAnnotations: List<FirAnnotation>) {
         annotations = newAnnotations.toMutableOrEmpty()
     }
+
+    override fun replaceBlock(newBlock: FirBlock) {
+        block = newBlock
+    }
 }
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt
index dfe1f22fa58..7b8a1aebe12 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/BuilderConfigurator.kt
@@ -370,6 +370,10 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             withCopy()
         }
 
+        builder(codeFragment) {
+            withCopy()
+        }
+
         builder(resolvedQualifier) {
             parents += abstractResolvedQualifierBuilder
             defaultFalse("isNullableLHSForCallableReference", "isFullyQualified")
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 a98b2fdac30..3d7dedba2ab 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
@@ -497,7 +497,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
 
         codeFragment.configure {
             +symbol("FirCodeFragmentSymbol")
-            +field(block).withTransform()
+            +field(block, withReplace = true).withTransform()
         }
 
         packageDirective.configure {