[FIR/IR generator] Commonize visitor printing logic

This is a step towards commonizing the code generator between FIR and IR: KT-61970 Also, don't use kotlinpoet for generating IR visitors (KT-61703)
2023-09-20 14:13:40 +02:00
parent 205a125c5f
commit c5f519f7c7
26 changed files with 2708 additions and 2393 deletions
@@ -734,6 +734,10 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
            +field("partiallyResolvedTypeRef", typeRef, nullable = true).withTransform()
        }

+        resolvedErrorReference.configure {
+            element.customParentInVisitor = resolvedNamedReference
+        }
+
        intersectionTypeRef.configure {
            +field("leftType", typeRef)
            +field("rightType", typeRef)
@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.fir.tree.generator.context.generatedType
 import org.jetbrains.kotlin.fir.tree.generator.context.type
+import org.jetbrains.kotlin.fir.tree.generator.printer.VISITOR_PACKAGE
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -116,11 +117,11 @@ val annotationResolvePhaseType = generatedType("expressions", "FirAnnotationReso

 val typeRefMarkerType = type("mpp", "TypeRefMarker")

-val firVisitorType = type("fir.visitors", "FirVisitor", kind = TypeKind.Class)
-val firVisitorVoidType = type("fir.visitors", "FirVisitorVoid", kind = TypeKind.Class)
-val firDefaultVisitorType = type("fir.visitors", "FirDefaultVisitor", kind = TypeKind.Class)
-val firDefaultVisitorVoidType = type("fir.visitors", "FirDefaultVisitorVoid", kind = TypeKind.Class)
-val firTransformerType = type("fir.visitors", "FirTransformer", kind = TypeKind.Class)
+val firVisitorType = generatedType("visitors", "FirVisitor")
+val firVisitorVoidType = generatedType("visitors", "FirVisitorVoid")
+val firDefaultVisitorType = generatedType("visitors", "FirDefaultVisitor")
+val firDefaultVisitorVoidType = generatedType("visitors", "FirDefaultVisitorVoid")
+val firTransformerType = generatedType("visitors", "FirTransformer")

 val resolveStateAccessAnnotation = type("fir.declarations", "ResolveStateAccess", kind = TypeKind.Class)
 val unresolvedExpressionTypeAccessAnnotation = type("fir.expressions", "UnresolvedExpressionTypeAccess", kind = TypeKind.Class)
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.tree.generator.model

 import org.jetbrains.kotlin.fir.tree.generator.printer.BASE_PACKAGE
+import org.jetbrains.kotlin.fir.tree.generator.printer.safeDecapitalizedName
 import org.jetbrains.kotlin.fir.tree.generator.util.set
 import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.ElementOrRef as GenericElementOrRef
@@ -78,6 +79,17 @@ class Element(override val name: String, override val propertyName: String, kind
    var baseTransformerType: Element? = null
    val transformerType: Element get() = baseTransformerType ?: this

+    override val visitFunctionName: String
+        get() = "visit$name"
+
+    override val visitorParameterName: String
+        get() = safeDecapitalizedName
+
+    var customParentInVisitor: Element? = null
+
+    override val parentInVisitor: Element?
+        get() = customParentInVisitor ?: elementParents.singleOrNull()?.element?.takeIf { !it.isRootElement }
+
    var doesNotNeedImplementation: Boolean = false

    val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || elementParents.any { it.element.needTransformOtherChildren }
@@ -40,7 +40,7 @@ fun SmartPrinter.printElement(element: Element) {

        printKDoc(element.extendedKDoc())
        print("${kind!!.title} $typeName")
-        print(typeParameters())
+        print(params.typeParameters())
        val parentRefs = element.parentRefs
        if (parentRefs.isNotEmpty()) {
            print(
@@ -49,7 +49,7 @@ fun SmartPrinter.printElement(element: Element) {
                }
            )
        }
-        print(multipleUpperBoundsList())
+        print(params.multipleUpperBoundsList())
        println(" {")
        withIndent {
            allFields.forEach { field ->
@@ -67,7 +67,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
            print("internal ")
        }
        print("${kind!!.title} ${this.typeName}")
-        print(element.typeParameters(end = " "))
+        print(element.params.typeParameters(end = " "))

        val isInterface = kind == ImplementationKind.Interface || kind == ImplementationKind.SealedInterface
        val isAbstract = kind == ImplementationKind.AbstractClass || kind == ImplementationKind.SealedClass
@@ -5,71 +5,107 @@

 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.firTransformerType
+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.GeneratedFile
-import org.jetbrains.kotlin.generators.tree.printer.multipleUpperBoundsList
+import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
 import org.jetbrains.kotlin.generators.tree.printer.printGeneratedType
-import org.jetbrains.kotlin.generators.tree.printer.typeParameters
-import org.jetbrains.kotlin.generators.tree.render
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
 import java.io.File

-fun printTransformer(elements: List<Element>, generationPath: File): GeneratedFile =
-    printGeneratedType(generationPath, TREE_GENERATOR_README, VISITOR_PACKAGE, "FirTransformer") {
-        println()
-        println("abstract class FirTransformer<in D> : FirVisitor<FirElement, D>() {")
-        println()
-        withIndent {
-            println("abstract fun <E : FirElement> transformElement(element: E, data: D): E")
+private class TransformerPrinter(
+    printer: SmartPrinter,
+) : AbstractVisitorPrinter<Element, Field>(printer, visitSuperTypeByDefault = false) {
+
+    override val visitorType: ClassRef<*>
+        get() = firTransformerType
+
+    override val visitorSuperType: ClassRef<PositionTypeParameterRef>
+        get() = firVisitorType.withArgs(AbstractFirTreeBuilder.baseFirElement, visitorDataType)
+
+    override val visitorTypeParameters: List<TypeVariable>
+        get() = listOf(dataTypeVariable)
+
+    override val visitorDataType: TypeRef
+        get() = dataTypeVariable
+
+    override fun visitMethodReturnType(element: Element) = element.transformerType
+
+    override val allowTypeParametersInVisitorMethods: Boolean
+        get() = true
+
+    context(ImportCollector)
+    override fun printMethodsForElement(element: Element) {
+        printer.run {
            println()
-            for (element in elements) {
-                if (element == AbstractFirTreeBuilder.baseFirElement) continue
-                val varName = element.safeDecapitalizedName
-                print("open fun ")
-                element.typeParameters(end = " ").takeIf { it.isNotBlank() }?.let { print(it) }
-                println(
-                    "transform",
-                    element.name,
-                    "(",
-                    varName,
-                    ": ",
-                    element.render(),
-                    ", data: D): ",
-                    element.transformerType.render(),
-                    element.multipleUpperBoundsList(),
-                    " {",
+            val elementParameterName = element.safeDecapitalizedName
+            if (element.isRootElement) {
+                val elementTP = TypeVariable("E", listOf(element), variance = Variance.INVARIANT)
+                printFunctionDeclaration(
+                    name = "transformElement",
+                    parameters = listOf(
+                        FunctionParameter(elementParameterName, elementTP),
+                        FunctionParameter("data", dataTypeVariable)
+                    ),
+                    returnType = elementTP,
+                    typeParameters = listOf(elementTP),
+                    modality = Modality.ABSTRACT,
                )
+                println()
+            } else {
+                printFunctionDeclaration(
+                    name = "transform" + element.name,
+                    parameters = listOf(
+                        FunctionParameter(elementParameterName, element),
+                        FunctionParameter("data", dataTypeVariable)
+                    ),
+                    returnType = visitMethodReturnType(element),
+                    typeParameters = element.params,
+                    modality = Modality.OPEN,
+                )
+                println(" {")
                withIndent {
-                    println("return transformElement($varName, data)")
+                    println("return transformElement(", elementParameterName, ", data)")
                }
                println("}")
-                println()
            }
-
-            for (element in elements) {
-                val varName = element.safeDecapitalizedName
-                print("final override fun ")
-                element.typeParameters(end = " ").takeIf { it.isNotBlank() }?.let { print(it) }
-
+            println()
+            printVisitMethodDeclaration(
+                element = element,
+                modality = Modality.FINAL,
+                override = true,
+            )
+            println(" {")
+            withIndent {
                println(
-                    "visit",
+                    "return transform",
                    element.name,
                    "(",
-                    varName,
-                    ": ",
-                    element.render(),
-                    ", data: D): ",
-                    element.transformerType.render(),
-                    element.multipleUpperBoundsList(),
-                    " {"
+                    element.safeDecapitalizedName,
+                    ", ",
+                    "data)"
                )
-                withIndent {
-                    println("return transform${element.name}($varName, data)")
-                }
-                println("}")
-                println()
            }
+            println("}")
        }
-        println("}")
+    }
+}
+
+fun printTransformer(elements: List<Element>, generationPath: File): GeneratedFile =
+    printGeneratedType(
+        generationPath,
+        TREE_GENERATOR_README,
+        firTransformerType.packageName,
+        firTransformerType.simpleName,
+    ) {
+        println()
+        TransformerPrinter(this).printVisitor(elements)
    }
@@ -5,176 +5,121 @@

 package org.jetbrains.kotlin.fir.tree.generator.printer

-import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
+import org.jetbrains.kotlin.fir.tree.generator.firDefaultVisitorType
+import org.jetbrains.kotlin.fir.tree.generator.firDefaultVisitorVoidType
+import org.jetbrains.kotlin.fir.tree.generator.firVisitorType
+import org.jetbrains.kotlin.fir.tree.generator.firVisitorVoidType
 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.GeneratedFile
-import org.jetbrains.kotlin.generators.tree.printer.multipleUpperBoundsList
 import org.jetbrains.kotlin.generators.tree.printer.printGeneratedType
-import org.jetbrains.kotlin.generators.tree.printer.typeParameters
-import org.jetbrains.kotlin.generators.tree.render
-import org.jetbrains.kotlin.utils.withIndent
+import org.jetbrains.kotlin.utils.SmartPrinter
 import java.io.File

-private val elementsWithMultipleSupertypesForDefaultVisitor = mapOf(
-    FirTreeBuilder.resolvedErrorReference to FirTreeBuilder.resolvedNamedReference
-)
+private class VisitorPrinter(
+    printer: SmartPrinter,
+    override val visitorType: ClassRef<*>,
+    visitSuperTypeByDefault: Boolean,
+) : AbstractVisitorPrinter<Element, Field>(printer, visitSuperTypeByDefault) {

-private fun Element.isAcceptableForDefaultVisiting(): Boolean {
-    if (this == AbstractFirTreeBuilder.baseFirElement) return false
-    val hasSingleSupertype = elementParents.size == 1 && !elementParents[0].element.isRootElement
-    return hasSingleSupertype || this in elementsWithMultipleSupertypesForDefaultVisitor
-}
+    override val visitorTypeParameters: List<TypeVariable>
+        get() = listOf(resultTypeVariable, dataTypeVariable)

-private fun Element.getNameOfSupertypeForDefaultVisiting(): String {
-    val parentForDefaultVisiting =
-        elementParents.singleOrNull()?.element ?: elementsWithMultipleSupertypesForDefaultVisitor.getValue(this)
-    return parentForDefaultVisiting.name
-}
+    override val visitorSuperType: ClassRef<PositionTypeParameterRef>? =
+        firVisitorType.takeIf { visitSuperTypeByDefault }?.withArgs(resultTypeVariable, dataTypeVariable)

-fun printVisitor(elements: List<Element>, generationPath: File, visitSuperTypeByDefault: Boolean): GeneratedFile {
-    val className = if (visitSuperTypeByDefault) "FirDefaultVisitor" else "FirVisitor"
-    return printGeneratedType(generationPath, TREE_GENERATOR_README, VISITOR_PACKAGE, className) {
-        println()
+    override val visitorDataType: TypeRef
+        get() = dataTypeVariable

-        print("abstract class $className<out R, in D> ")
-        if (visitSuperTypeByDefault) {
-            print(": FirVisitor<R, D>() ")
-        }
-        println("{")
+    override fun visitMethodReturnType(element: Element) = resultTypeVariable

-        pushIndent()
-        if (!visitSuperTypeByDefault) {
-            println("abstract fun visitElement(element: FirElement, data: D): R\n")
-        }
-        for (element in elements) {
-            if (element == AbstractFirTreeBuilder.baseFirElement) continue
-            if (visitSuperTypeByDefault && !element.isAcceptableForDefaultVisiting()) continue
-            with(element) {
-                val varName = safeDecapitalizedName
-                if (visitSuperTypeByDefault) {
-                    print("override")
-                } else {
-                    print("open")
-                }
-                print(
-                    " fun ",
-                    typeParameters(end = " "),
-                    "visit",
-                    name,
-                    "(",
-                    varName,
-                    ": ",
-                    render(),
-                    ", data: D): R",
-                    multipleUpperBoundsList(),
-                    " = visit"
-                )
-                if (visitSuperTypeByDefault) {
-                    print(element.getNameOfSupertypeForDefaultVisiting())
-                } else {
-                    print("Element")
-                }
-                println("($varName, data)")
-                println()
-            }
-        }
-        popIndent()
-        println("}")
+    override val allowTypeParametersInVisitorMethods: Boolean
+        get() = true
+
+    override fun parentInVisitor(element: Element): Element? = when {
+        element.isRootElement -> null
+        visitSuperTypeByDefault -> element.parentInVisitor
+        else -> AbstractFirTreeBuilder.baseFirElement
    }
 }

+fun printVisitor(elements: List<Element>, generationPath: File, visitSuperTypeByDefault: Boolean) =
+    printVisitorCommon(
+        elements,
+        generationPath,
+        if (visitSuperTypeByDefault) firDefaultVisitorType else firVisitorType,
+    ) { printer, visitorType ->
+        VisitorPrinter(printer, visitorType, visitSuperTypeByDefault)
+    }

-fun printVisitorVoid(elements: List<Element>, generationPath: File): GeneratedFile =
-    printGeneratedType(generationPath, TREE_GENERATOR_README, VISITOR_PACKAGE, "FirVisitorVoid") {
-        println()
-        println("abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {")
+private class VisitorVoidPrinter(
+    printer: SmartPrinter,
+    override val visitorType: ClassRef<*>,
+) : AbstractVisitorVoidPrinter<Element, Field>(printer, visitSuperTypeByDefault = false) {

-        withIndent {
-            println("abstract fun visitElement(element: FirElement)")
+    override val visitorSuperClass: ClassRef<PositionTypeParameterRef>
+        get() = firVisitorType
+
+    override val allowTypeParametersInVisitorMethods: Boolean
+        get() = true
+
+    override val useAbstractMethodForRootElement: Boolean
+        get() = true
+
+    override val overriddenVisitMethodsAreFinal: Boolean
+        get() = true
+
+    override fun parentInVisitor(element: Element): Element = AbstractFirTreeBuilder.baseFirElement
+}
+
+fun printVisitorVoid(elements: List<Element>, generationPath: File) =
+    printVisitorCommon(elements, generationPath, firVisitorVoidType, ::VisitorVoidPrinter)
+
+private class DefaultVisitorVoidPrinter(
+    printer: SmartPrinter,
+    override val visitorType: ClassRef<*>,
+) : AbstractVisitorPrinter<Element, Field>(printer, visitSuperTypeByDefault = true) {
+
+    override val visitorTypeParameters: List<TypeVariable>
+        get() = emptyList()
+
+    override val visitorDataType: TypeRef
+        get() = StandardTypes.nothing.copy(nullable = true)
+
+    override fun visitMethodReturnType(element: Element) = StandardTypes.unit
+
+    override val visitorSuperType: ClassRef<PositionTypeParameterRef>
+        get() = firVisitorVoidType
+
+    override val allowTypeParametersInVisitorMethods: Boolean
+        get() = true
+
+    context(ImportCollector)
+    override fun printMethodsForElement(element: Element) {
+        printer.run {
+            printVisitMethodDeclaration(
+                element,
+                hasDataParameter = false,
+                override = true,
+            )
+            println(" = ", element.parentInVisitor!!.visitFunctionName, "(", element.visitorParameterName, ")")
            println()
-            for (element in elements) {
-                if (element == AbstractFirTreeBuilder.baseFirElement) continue
-                with(element) {
-                    val varName = safeDecapitalizedName
-                    println(
-                        "open fun ",
-                        typeParameters(end = " "),
-                        "visit",
-                        name,
-                        "(",
-                        varName,
-                        ": ",
-                        render(),
-                        ")",
-                        multipleUpperBoundsList(),
-                        " {"
-                    )
-                    withIndent {
-                        println("visitElement($varName)")
-                    }
-                    println("}")
-                    println()
-                }
-            }
-
-            for (element in elements) {
-                with(element) {
-                    val varName = safeDecapitalizedName
-                    println(
-                        "final override fun ",
-                        typeParameters(end = " "),
-                        "visit",
-                        name,
-                        "(",
-                        varName,
-                        ": ",
-                        render(),
-                        ", data: Nothing?)",
-                        multipleUpperBoundsList(),
-                        " {"
-                    )
-                    withIndent {
-                        println("visit$name($varName)")
-                    }
-                    println("}")
-                    println()
-                }
-            }
        }
-        println("}")
    }
+}

-fun printDefaultVisitorVoid(elements: List<Element>, generationPath: File): GeneratedFile =
-    printGeneratedType(generationPath, TREE_GENERATOR_README, VISITOR_PACKAGE, "FirDefaultVisitorVoid") {
+fun printDefaultVisitorVoid(elements: List<Element>, generationPath: File) =
+    printVisitorCommon(elements, generationPath, firDefaultVisitorVoidType, ::DefaultVisitorVoidPrinter)
+
+private fun printVisitorCommon(
+    elements: List<Element>,
+    generationPath: File,
+    visitorType: ClassRef<*>,
+    makePrinter: (SmartPrinter, ClassRef<*>) -> AbstractVisitorPrinter<Element, Field>,
+): GeneratedFile =
+    printGeneratedType(generationPath, TREE_GENERATOR_README, visitorType.packageName, visitorType.simpleName) {
        println()
-        println("abstract class FirDefaultVisitorVoid : FirVisitorVoid() {")
-
-        pushIndent()
-        for (element in elements) {
-            if (!element.isAcceptableForDefaultVisiting()) continue
-            with(element) {
-                val varName = safeDecapitalizedName
-                println(
-                    "override fun ",
-                    typeParameters(end = " "),
-                    "visit",
-                    name,
-                    "(",
-                    varName,
-                    ": ",
-                    render(),
-                    ")",
-                    multipleUpperBoundsList(),
-                    " = visit",
-                    getNameOfSupertypeForDefaultVisiting(),
-                    "(",
-                    varName,
-                    ")"
-                )
-                println()
-            }
-        }
-        popIndent()
-        println("}")
-    }
+        makePrinter(this, visitorType).printVisitor(elements)
+    }