From b92194229c0bd98635d4895df7b5daa00ae37280 Mon Sep 17 00:00:00 2001
From: "simon.ogorodnik" <simon.ogorodnik@jetbrains.com>
Date: Wed, 20 May 2020 20:04:39 +0300
Subject: [PATCH] [FIR] Add fir copy function generation

---
 .../context/AbstractBuilderConfigurator.kt    |  4 ++
 .../fir/tree/generator/model/Builder.kt       |  1 +
 .../fir/tree/generator/printer/builder.kt     | 41 +++++++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractBuilderConfigurator.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractBuilderConfigurator.kt
index 3f7aa88f0e4..20f47ce40af 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractBuilderConfigurator.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/context/AbstractBuilderConfigurator.kt
@@ -152,6 +152,10 @@ abstract class AbstractBuilderConfigurator<T : AbstractFirTreeBuilder>(val firTr
         fun openBuilder() {
             builder.isOpen = true
         }
+
+        fun withCopy() {
+            builder.wantsCopy = true
+        }
     }
 
     fun builder(name: String? = null, block: IntermediateBuilderConfigurationContext.() -> Unit): IntermediateBuilderDelegateProvider {
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Builder.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Builder.kt
index f76d3e7492c..02f6304c109 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Builder.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Builder.kt
@@ -49,6 +49,7 @@ class LeafBuilder(val implementation: Implementation) : Builder() {
 
     override val packageName: String = implementation.packageName.replace(".impl", ".builder")
     var isOpen: Boolean = false
+    var wantsCopy: Boolean = false
 }
 
 class IntermediateBuilder(override val type: String) : Builder() {
diff --git a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/builder.kt b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/builder.kt
index e4389c1d0c3..e6050cb2df8 100644
--- a/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/builder.kt
+++ b/compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/builder.kt
@@ -111,6 +111,11 @@ private fun SmartPrinter.printBuilder(builder: Builder) {
     if (builder is LeafBuilder) {
         println()
         printDslBuildFunction(builder, hasRequiredFields)
+
+        if (builder.wantsCopy) {
+            println()
+            printDslBuildCopyFunction(builder, hasRequiredFields)
+        }
     }
 }
 
@@ -261,4 +266,40 @@ private fun SmartPrinter.printDslBuildFunction(
         }
     }
     println("}")
+}
+
+private fun SmartPrinter.printDslBuildCopyFunction(
+    builder: LeafBuilder,
+    hasRequiredFields: Boolean
+) {
+    println("@OptIn(ExperimentalContracts::class)")
+    print("inline ")
+    print("fun ")
+    builder.implementation.element.typeArguments.takeIf { it.isNotEmpty() }?.let {
+        print(it.joinToString(separator = ", ", prefix = "<", postfix = "> ") { it.name })
+    }
+    val builderType = builder.typeWithArguments
+    val name = builder.implementation.name?.replaceFirst("Fir", "") ?: builder.implementation.element.name
+    print("build${name}Copy(")
+    print("original: ${builder.implementation.element.typeWithArguments}, init: $builderType.() -> Unit")
+    if (!hasRequiredFields) {
+        print(" = {}")
+    }
+    println("): ${builder.implementation.element.typeWithArguments} {")
+    withIndent {
+        println("contract {")
+        withIndent {
+            println("callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)")
+        }
+        println("}")
+        println("val copyBuilder = $builderType()")
+        for (field in builder.allFields) {
+            when (field.origin) {
+                is FieldList -> println("copyBuilder.${field.name}.addAll(original.${field.name})")
+                else -> println("copyBuilder.${field.name} = original.${field.name}")
+            }
+        }
+        println("return copyBuilder.apply(init).build()")
+    }
+    println("}")
 }
\ No newline at end of file