[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)

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/TransformerPrinter.kt

@@ -17,7 +17,7 @@ internal class TransformerPrinter(
     printer: SmartPrinter,
     override val visitorType: ClassRef<*>,
     private val rootElement: Element,
-) : AbstractVisitorPrinter<Element, Field>(printer) {
+) : AbstractTransformerPrinter<Element, Field>(printer) {
 
     override val visitorSuperType: ClassRef<PositionTypeParameterRef>
         get() = elementVisitorType.withArgs(rootElement, dataTypeVariable)
@@ -30,9 +30,6 @@ internal class TransformerPrinter(
 
     override fun visitMethodReturnType(element: Element) = element.getTransformExplicitType()
 
-    override val allowTypeParametersInVisitorMethods: Boolean
-        get() = false
-
     context(ImportCollector)
     override fun printMethodsForElement(element: Element) {
         printer.run {

compiler/ir/ir.tree/tree-generator/src/org/jetbrains/kotlin/ir/generator/print/TransformerVoidPrinter.kt

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.utils.withIndent
 internal class TransformerVoidPrinter(
     printer: SmartPrinter,
     override val visitorType: ClassRef<*>,
-) : AbstractVisitorPrinter<Element, Field>(printer) {
+) : AbstractTransformerPrinter<Element, Field>(printer) {
 
     override val visitorTypeParameters: List<TypeVariable>
         get() = emptyList()

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractTransformerPrinter.kt

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.generators.tree
+
+import org.jetbrains.kotlin.descriptors.Modality
+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
+
+abstract class AbstractTransformerPrinter<Element : AbstractElement<Element, Field, *>, Field : AbstractField<Field>>(
+    printer: SmartPrinter,
+) : AbstractVisitorPrinter<Element, Field>(printer) {
+
+    override fun visitMethodReturnType(element: Element) = element.transformerClass
+
+    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 {
+                val parentInVisitor = parentInVisitor(element) ?: return
+                printFunctionWithBlockBody(
+                    name = "transform" + element.name,
+                    parameters = listOf(
+                        FunctionParameter(elementParameterName, element),
+                        FunctionParameter("data", dataTypeVariable)
+                    ),
+                    returnType = visitMethodReturnType(element),
+                    typeParameters = element.params,
+                    modality = Modality.OPEN,
+                ) {
+                    println("return transform", parentInVisitor.name, "(", elementParameterName, ", data)")
+                }
+            }
+            println()
+            printVisitMethodDeclaration(
+                element = element,
+                modality = Modality.FINAL,
+                override = true,
+            )
+            printBlock {
+                println(
+                    "return transform",
+                    element.name,
+                    "(",
+                    element.visitorParameterName,
+                    ", ",
+                    "data)"
+                )
+            }
+        }
+    }
+}

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/utils.kt

@@ -79,3 +79,9 @@ internal fun <Field : AbstractField<*>> List<Field>.reorderFieldsIfNecessary(ord
 
 val ImplementationKind.hasLeafBuilder: Boolean
     get() = this == ImplementationKind.FinalClass || this == ImplementationKind.OpenClass
+
+val AbstractElement<*, *, *>.safeDecapitalizedName: String
+    get() = when (name) {
+        "Class" -> "klass"
+        else -> name.replaceFirstChar(Char::lowercaseChar)
+    }

native/swift/sir/tree-generator/src/org/jetbrains/kotlin/sir/tree/generator/model/Element.kt

@@ -20,10 +20,8 @@ class Element(name: String, override val propertyName: String) : AbstractElement
 
     override val otherParents: MutableList<ClassRef<*>> = mutableListOf()
 
-    override val visitorParameterName: String = when (val parameterName = name.replaceFirstChar(Char::lowercaseChar)) {
-        "class" -> "klass"
-        else -> parameterName
-    }
+    override val visitorParameterName: String
+        get() = safeDecapitalizedName
 
     override val hasAcceptMethod: Boolean
         get() = true

native/swift/sir/tree-generator/src/org/jetbrains/kotlin/sir/tree/generator/printer/TransformerPrinter.kt

@@ -5,12 +5,7 @@
 
 package org.jetbrains.kotlin.sir.tree.generator.printer
 
-import org.jetbrains.kotlin.descriptors.Modality
 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.sir.tree.generator.elementVisitorType
 import org.jetbrains.kotlin.sir.tree.generator.model.Element
 import org.jetbrains.kotlin.sir.tree.generator.model.Field
@@ -20,7 +15,7 @@ internal class TransformerPrinter(
     printer: SmartPrinter,
     override val visitorType: ClassRef<*>,
     private val rootElement: Element,
-) : AbstractVisitorPrinter<Element, Field>(printer) {
+) : AbstractTransformerPrinter<Element, Field>(printer) {
 
     override val visitorSuperType: ClassRef<PositionTypeParameterRef>
         get() = elementVisitorType.withArgs(rootElement, dataTypeVariable)
@@ -30,59 +25,4 @@ internal class TransformerPrinter(
 
     override val visitorDataType: TypeRef
         get() = dataTypeVariable
-
-    override fun visitMethodReturnType(element: Element) = element.transformerClass
-
-    context(ImportCollector)
-    override fun printMethodsForElement(element: Element) {
-        // FIXME: This code is copy-pasted from the FIR generator. Factor it out.
-        printer.run {
-            println()
-            val elementParameterName = element.visitorParameterName
-            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 {
-                val parentInVisitor = element.parentInVisitor ?: return
-                printFunctionWithBlockBody(
-                    name = "transform" + element.name,
-                    parameters = listOf(
-                        FunctionParameter(elementParameterName, element),
-                        FunctionParameter("data", dataTypeVariable)
-                    ),
-                    returnType = visitMethodReturnType(element),
-                    typeParameters = element.params,
-                    modality = Modality.OPEN,
-                ) {
-                    println("return ", "transform", parentInVisitor.name, "(", elementParameterName, ", data)")
-                }
-            }
-            println()
-            printVisitMethodDeclaration(
-                element = element,
-                modality = Modality.FINAL,
-                override = true,
-            )
-            printBlock {
-                println(
-                    "return transform",
-                    element.name,
-                    "(",
-                    elementParameterName,
-                    ", ",
-                    "data)"
-                )
-            }
-        }
-    }
 }