[FIR] Add fir copy function generation

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 {

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() {

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("}")
 }