[FIR generator] Factor out element printer to common module

This is a step towards commonizing the code generator between
FIR and IR: KT-61970

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

@@ -77,7 +77,7 @@ class Element(override val name: String, override val propertyName: String, kind
         get() = emptyList() // Use Implementation#transformableChildren instead
 
     var baseTransformerType: Element? = null
-    val transformerType: Element get() = baseTransformerType ?: this
+    val transformerClass: Element get() = baseTransformerType ?: this
 
     override val visitFunctionName: String
         get() = "visit$name"

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

@@ -13,140 +13,81 @@ import org.jetbrains.kotlin.fir.tree.generator.util.get
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.printer.*
 import org.jetbrains.kotlin.utils.SmartPrinter
-import org.jetbrains.kotlin.utils.withIndent
 import java.io.File
 
-fun Element.generateCode(generationPath: File): GeneratedFile =
-    printGeneratedType(generationPath, TREE_GENERATOR_README, packageName, this.typeName) {
-        println()
-        printElement(this@generateCode)
-    }
+private class ElementPrinter(printer: SmartPrinter) : AbstractElementPrinter<Element, Field>(printer) {
 
-private class ElementFieldPrinter(printer: SmartPrinter) : AbstractFieldPrinter<Field>(printer)
+    override fun makeFieldPrinter(printer: SmartPrinter) = object : AbstractFieldPrinter<Field>(printer) {}
 
-context(ImportCollector)
-fun SmartPrinter.printElement(element: Element) {
-    with(element) {
-        val isInterface = kind == ImplementationKind.Interface || kind == ImplementationKind.SealedInterface
-        fun abstract() {
-            if (!isInterface) {
-                print("abstract ")
-            }
-        }
+    context(ImportCollector)
+    override fun SmartPrinter.printAdditionalMethods(element: Element) {
+        val kind = element.kind ?: error("Expected non-null element kind")
+        with(element) {
+            // TODO: Add a kDoc for `accept`
+            printAcceptMethod(element, firVisitorType, hasImplementation = true, kDoc = null)
 
-        fun override() {
-            if (!isRootElement) {
-                print("override ")
-            }
-        }
-
-        printKDoc(element.extendedKDoc())
-        print("${kind!!.title} $typeName")
-        print(params.typeParameters())
-        val parentRefs = element.parentRefs
-        if (parentRefs.isNotEmpty()) {
-            print(
-                parentRefs.sortedBy { it.typeKind }.joinToString(prefix = " : ") { parent ->
-                    parent.render() + parent.inheritanceClauseParenthesis()
-                }
+            // TODO: Add a kDoc for `transform`
+            printTransformMethod(
+                element = element,
+                transformerClass = firTransformerType,
+                implementation = "transformer.transform${element.name}(this, data)",
+                returnType = TypeVariable("E", listOf(AbstractFirTreeBuilder.baseFirElement)),
+                kDoc = null,
             )
-        }
-        print(params.multipleUpperBoundsList())
-        println(" {")
-        withIndent {
-            val fieldPrinter = ElementFieldPrinter(this@printElement)
-            allFields.forEach { field ->
-                if (field.isFinal && field.fromParent || field.isParameter) return@forEach
-                fieldPrinter.printField(field, override = field.fromParent) {
-                    if (!field.isFinal) {
-                        abstract()
-                    }
-                }
-            }
-
-            if (hasAcceptMethod) {
-                if (allFields.isNotEmpty()) {
-                    println()
-                }
-
-                override()
-                println("fun <R, D> accept(visitor: ${firVisitorType.render()}<R, D>, data: D): R =")
-                withIndent {
-                    println("visitor.visit${element.name}(this, data)")
-                }
-            }
-
-            if (hasTransformMethod) {
-                println()
-                println("@Suppress(\"UNCHECKED_CAST\")")
-                override()
-                println(
-                    "fun <E : ",
-                    AbstractFirTreeBuilder.baseFirElement.render(),
-                    ", D> transform(transformer: ",
-                    firTransformerType.render(),
-                    "<D>, data: D): E ="
-                )
-                withIndent {
-                    println("transformer.transform$name(this, data) as E")
-                }
-            }
 
             fun Field.replaceDeclaration(override: Boolean, overridenType: TypeRefWithNullability? = null, forceNullable: Boolean = false) {
                 println()
                 if (name == "source") {
-                    println("@${firImplementationDetailType.render()}")
+                    println("@", firImplementationDetailType.render())
                 }
-                abstract()
-                if (override) print("override ")
-                println(replaceFunctionDeclaration(overridenType, forceNullable))
+                replaceFunctionDeclaration(this, override, kind, overridenType, forceNullable)
+                println()
             }
 
             allFields.filter { it.withReplace }.forEach {
-                val override = overridenFields[it, it] &&
-                        !(it.name == "source" && element == FirTreeBuilder.qualifiedAccessExpression)
+                val override = overridenFields[it, it] && !(it.name == "source" && element == FirTreeBuilder.qualifiedAccessExpression)
                 it.replaceDeclaration(override, forceNullable = it.useNullableForReplace)
-                for (overridenType in it.overridenTypes) {
-                    it.replaceDeclaration(true, overridenType)
+                for (overriddenType in it.overridenTypes) {
+                    it.replaceDeclaration(true, overriddenType)
                 }
             }
 
             for (field in allFields) {
                 if (!field.needsSeparateTransform) continue
                 println()
-                abstract()
-                if (field.fromParent && field.parentHasSeparateTransform) {
-                    print("override ")
-                }
-                println(field.transformFunctionDeclaration(element))
+                transformFunctionDeclaration(field, element, override = field.fromParent && field.parentHasSeparateTransform, kind)
+                println()
             }
             if (needTransformOtherChildren) {
                 println()
-                abstract()
-                if (element.elementParents.any { it.element.needTransformOtherChildren }) {
-                    print("override ")
-                }
-                println(transformFunctionDeclaration("OtherChildren", element))
+                transformOtherChildrenFunctionDeclaration(
+                    element,
+                    override = element.elementParents.any { it.element.needTransformOtherChildren },
+                    kind,
+                )
+                println()
             }
 
             if (element.isRootElement) {
-                require(isInterface) {
-                    "$element must be an interface"
-                }
                 println()
-                println("fun accept(visitor: ${firVisitorVoidType.render()}) = accept(visitor, null)")
-                if (element.hasAcceptChildrenMethod) {
-                    println()
-                    println("fun <R, D> acceptChildren(visitor: ${firVisitorType.render()}<R, D>, data: D)")
-                }
+                println("fun accept(visitor: ", firVisitorVoidType.render(), ") = accept(visitor, null)")
+
+                // TODO: Add a kDoc for `acceptChildren`
+                printAcceptChildrenMethod(element, firVisitorType, visitorResultType = TypeVariable("R"), kDoc = null)
+                println()
+                println()
+                println("fun acceptChildren(visitor: ", firVisitorVoidType.render(), ") = acceptChildren(visitor, null)")
+
+                // TODO: Add a kDoc for `transformChildren`
+                printTransformChildrenMethod(element, firTransformerType, returnType = AbstractFirTreeBuilder.baseFirElement, kDoc = null)
                 println()
-                println("fun acceptChildren(visitor: ${firVisitorVoidType.render()}) = acceptChildren(visitor, null)")
-                if (element.hasTransformChildrenMethod) {
-                    println()
-                    println("fun <D> transformChildren(transformer: ${firTransformerType.render()}<D>, data: D): FirElement")
-                }
             }
         }
-        println("}")
     }
 }
+
+fun Element.generateCode(generationPath: File): GeneratedFile =
+    printGeneratedType(generationPath, TREE_GENERATOR_README, packageName, typeName) {
+        println()
+        ElementPrinter(this).printElement(element)
+    }

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

@@ -9,10 +9,8 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.tree.generator.*
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.generators.tree.printer.GeneratedFile
+import org.jetbrains.kotlin.generators.tree.printer.*
 import org.jetbrains.kotlin.generators.tree.printer.braces
-import org.jetbrains.kotlin.generators.tree.printer.printGeneratedType
-import org.jetbrains.kotlin.generators.tree.printer.typeParameters
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 import org.jetbrains.kotlin.utils.withIndent
@@ -86,12 +84,6 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
         val isInterface = kind == ImplementationKind.Interface || kind == ImplementationKind.SealedInterface
         val isAbstract = kind == ImplementationKind.AbstractClass || kind == ImplementationKind.SealedClass
 
-        fun abstract() {
-            if (isAbstract) {
-                print("abstract ")
-            }
-        }
-
         val fieldPrinter = ImplementationFieldPrinter(this@printImplementation)
 
         if (!isInterface && !isAbstract && fieldsWithoutDefault.isNotEmpty()) {
@@ -128,9 +120,6 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                 fieldsWithDefault.forEach {
                     fieldPrinter.printField(it, override = true)
                 }
-                if (fieldsWithDefault.isNotEmpty()) {
-                    println()
-                }
             }
 
 
@@ -146,6 +135,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
 
             val customCalls = fieldsWithoutDefault.filter { it.customInitializationCall != null }
             if (bindingCalls.isNotEmpty() || customCalls.isNotEmpty()) {
+                println()
                 println("init {")
                 withIndent {
                     for (symbolField in bindingCalls) {
@@ -158,12 +148,13 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                     }
                 }
                 println("}")
-                println()
             }
 
             fun Field.acceptString(): String = "${name}${call()}accept(visitor, data)"
+
             if (hasAcceptChildrenMethod) {
-                print("override fun <R, D> acceptChildren(visitor: ${firVisitorType.render()}<R, D>, data: D) {")
+                printAcceptChildrenMethod(this, firVisitorType, TypeVariable("R"), override = true, kDoc = null)
+                print(" {")
 
                 val walkableFields = walkableChildren
                 if (walkableFields.isNotEmpty()) {
@@ -223,16 +214,16 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                     }
                 }
                 println("}")
-                println()
             }
 
             if (hasTransformChildrenMethod) {
-                abstract()
-                print(
-                    "override fun <D> transformChildren(transformer: ",
-                    firTransformerType.render(),
-                    "<D>, data: D): ",
-                    render(),
+                printTransformChildrenMethod(
+                    this,
+                    firTransformerType,
+                    this,
+                    modality = Modality.ABSTRACT.takeIf { isAbstract },
+                    override = true,
+                    kDoc = null,
                 )
                 if (!isInterface && !isAbstract) {
                     println(" {")
@@ -286,17 +277,15 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                         }
                         println("return this")
                     }
-                    println("}")
-                } else {
-                    println()
+                    print("}")
                 }
+                println()
             }
 
             for (field in allFields) {
                 if (!field.needsSeparateTransform) continue
                 println()
-                abstract()
-                print("override ${field.transformFunctionDeclaration(this)}")
+                transformFunctionDeclaration(field, this, override = true, kind!!)
                 if (isInterface || isAbstract) {
                     println()
                     continue
@@ -327,13 +316,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
 
             if (element.needTransformOtherChildren) {
                 println()
-                abstract()
-                print(
-                    "override fun <D> transformOtherChildren(transformer: ",
-                    firTransformerType.render(),
-                    "<D>, data: D): ",
-                    render(),
-                )
+                transformOtherChildrenFunctionDeclaration(this, override = true, kind!!)
                 if (isInterface || isAbstract) {
                     println()
                 } else {
@@ -364,8 +347,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                 if (field.name == "source") {
                     println("@${firImplementationDetailType.render()}")
                 }
-                abstract()
-                print("override ${field.replaceFunctionDeclaration(overridenType, forceNullable)}")
+                replaceFunctionDeclaration(field, override = true, kind!!, overridenType, forceNullable)
                 if (isInterface || isAbstract) {
                     println()
                     return

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

@@ -37,7 +37,7 @@ private class TransformerPrinter(
     override val visitorDataType: TypeRef
         get() = dataTypeVariable
 
-    override fun visitMethodReturnType(element: Element) = element.transformerType
+    override fun visitMethodReturnType(element: Element) = element.transformerClass
 
     override val allowTypeParametersInVisitorMethods: Boolean
         get() = true

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

@@ -5,30 +5,77 @@
 
 package org.jetbrains.kotlin.fir.tree.generator.printer
 
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.tree.generator.firTransformerType
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.printer.FunctionParameter
+import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
+import org.jetbrains.kotlin.utils.SmartPrinter
 
 context(ImportCollector)
-fun Field.transformFunctionDeclaration(returnType: TypeRef): String {
-    return transformFunctionDeclaration(name.replaceFirstChar(Char::uppercaseChar), returnType)
+fun SmartPrinter.transformFunctionDeclaration(
+    field: Field,
+    returnType: TypeRef,
+    override: Boolean,
+    implementationKind: ImplementationKind,
+) {
+    transformFunctionDeclaration(field.name.replaceFirstChar(Char::uppercaseChar), returnType, override, implementationKind)
 }
 
 context(ImportCollector)
-fun transformFunctionDeclaration(transformName: String, returnType: TypeRef): String {
-    return "fun <D> transform$transformName(transformer: ${firTransformerType.render()}<D>, data: D): " +
-            returnType.render()
+fun SmartPrinter.transformOtherChildrenFunctionDeclaration(
+    element: TypeRef,
+    override: Boolean,
+    implementationKind: ImplementationKind,
+) {
+    transformFunctionDeclaration("OtherChildren", element, override, implementationKind)
 }
 
 context(ImportCollector)
-fun Field.replaceFunctionDeclaration(
-    overridenType: TypeRefWithNullability? = null,
+private fun SmartPrinter.transformFunctionDeclaration(
+    transformName: String,
+    returnType: TypeRef,
+    override: Boolean,
+    implementationKind: ImplementationKind,
+) {
+    val dataTP = TypeVariable("D")
+    printFunctionDeclaration(
+        name = "transform$transformName",
+        parameters = listOf(
+            FunctionParameter("transformer", firTransformerType.withArgs(dataTP)),
+            FunctionParameter("data", dataTP),
+        ),
+        returnType = returnType,
+        typeParameters = listOf(dataTP),
+        modality = Modality.ABSTRACT.takeIf {
+            implementationKind == ImplementationKind.AbstractClass || implementationKind == ImplementationKind.SealedClass
+        },
+        override = override,
+    )
+}
+
+context(ImportCollector)
+fun SmartPrinter.replaceFunctionDeclaration(
+    field: Field,
+    override: Boolean,
+    implementationKind: ImplementationKind,
+    overriddenType: TypeRefWithNullability? = null,
     forceNullable: Boolean = false,
-): String {
-    val capName = name.replaceFirstChar(Char::uppercaseChar)
-    val type = overridenType ?: typeRef
+) {
+    val capName = field.name.replaceFirstChar(Char::uppercaseChar)
+    val type = overriddenType ?: field.typeRef
     val typeWithNullable = if (forceNullable) type.copy(nullable = true) else type
-    return "fun replace$capName(new$capName: ${typeWithNullable.render()})"
+
+    printFunctionDeclaration(
+        name = "replace$capName",
+        parameters = listOf(FunctionParameter("new$capName", typeWithNullable)),
+        returnType = StandardTypes.unit,
+        modality = Modality.ABSTRACT.takeIf {
+            implementationKind == ImplementationKind.AbstractClass || implementationKind == ImplementationKind.SealedClass
+        },
+        override = override,
+    )
 }
 
 fun Field.getMutableType(forBuilder: Boolean = false): TypeRefWithNullability = when (this) {