[tree generators] Factor out common transformer printing logic

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

@@ -31,7 +31,7 @@ fun main(args: Array<String>) {
             firDefaultVisitorType to ::VisitorPrinter.bind(true),
             firVisitorVoidType to ::VisitorVoidPrinter,
             firDefaultVisitorVoidType to ::DefaultVisitorVoidPrinter,
-            firTransformerType to ::TransformerPrinter,
+            firTransformerType to ::TransformerPrinter.bind(model.rootElement),
         ),
         ImplementationConfigurator,
         BuilderConfigurator,

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.fir.tree.generator.BASE_PACKAGE
-import org.jetbrains.kotlin.fir.tree.generator.printer.safeDecapitalizedName
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
 import org.jetbrains.kotlin.generators.tree.ElementRef as GenericElementRef

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

@@ -5,22 +5,18 @@
 
 package org.jetbrains.kotlin.fir.tree.generator.printer
 
-import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.firVisitorType
 import org.jetbrains.kotlin.fir.tree.generator.model.Element
 import org.jetbrains.kotlin.fir.tree.generator.model.Field
 import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.generators.tree.printer.FunctionParameter
-import org.jetbrains.kotlin.generators.tree.printer.printBlock
-import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
-import org.jetbrains.kotlin.generators.tree.printer.printFunctionWithBlockBody
 import org.jetbrains.kotlin.utils.SmartPrinter
 
 internal class TransformerPrinter(
     printer: SmartPrinter,
     override val visitorType: ClassRef<*>,
-) : AbstractVisitorPrinter<Element, Field>(printer) {
+    private val rootElement: Element,
+) : AbstractTransformerPrinter<Element, Field>(printer) {
 
     override val visitorSuperType: ClassRef<PositionTypeParameterRef>
         get() = firVisitorType.withArgs(AbstractFirTreeBuilder.baseFirElement, visitorDataType)
@@ -31,59 +27,11 @@ internal class TransformerPrinter(
     override val visitorDataType: TypeRef
         get() = dataTypeVariable
 
-    override fun visitMethodReturnType(element: Element) = element.transformerClass
-
     override val allowTypeParametersInVisitorMethods: Boolean
         get() = true
 
-    context(ImportCollector)
-    override fun printMethodsForElement(element: Element) {
-        printer.run {
-            println()
-            val elementParameterName = element.safeDecapitalizedName
-            if (element.isRootElement) {
-                val elementTP = TypeVariable("E", listOf(element))
-                printFunctionDeclaration(
-                    name = "transformElement",
-                    parameters = listOf(
-                        FunctionParameter(elementParameterName, elementTP),
-                        FunctionParameter("data", dataTypeVariable)
-                    ),
-                    returnType = elementTP,
-                    typeParameters = listOf(elementTP),
-                    modality = Modality.ABSTRACT,
-                )
-                println()
-            } else {
-                printFunctionWithBlockBody(
-                    name = "transform" + element.name,
-                    parameters = listOf(
-                        FunctionParameter(elementParameterName, element),
-                        FunctionParameter("data", dataTypeVariable)
-                    ),
-                    returnType = visitMethodReturnType(element),
-                    typeParameters = element.params,
-                    modality = Modality.OPEN,
-                ) {
-                    println("return transformElement(", elementParameterName, ", data)")
-                }
-            }
-            println()
-            printVisitMethodDeclaration(
-                element = element,
-                modality = Modality.FINAL,
-                override = true,
-            )
-            printBlock {
-                println(
-                    "return transform",
-                    element.name,
-                    "(",
-                    element.safeDecapitalizedName,
-                    ", ",
-                    "data)"
-                )
-            }
-        }
+    override fun parentInVisitor(element: Element) = when {
+        element.isRootElement -> null
+        else -> rootElement
     }
 }

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

@@ -8,7 +8,6 @@ package org.jetbrains.kotlin.fir.tree.generator.printer
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.tree.generator.BASE_PACKAGE
 import org.jetbrains.kotlin.fir.tree.generator.firTransformerType
-import org.jetbrains.kotlin.fir.tree.generator.model.Element
 import org.jetbrains.kotlin.fir.tree.generator.model.Field
 import org.jetbrains.kotlin.fir.tree.generator.model.FieldList
 import org.jetbrains.kotlin.fir.tree.generator.model.FieldWithDefault
@@ -91,5 +90,3 @@ fun Field.getMutableType(forBuilder: Boolean = false): TypeRefWithNullability =
     is FieldWithDefault -> if (isMutable) origin.getMutableType() else typeRef
     else -> typeRef
 }
-
-val Element.safeDecapitalizedName: String get() = if (name == "Class") "klass" else name.replaceFirstChar(Char::lowercaseChar)