[FIR generator] Factor out element printer to common module

This is a step towards commonizing the code generator between FIR and IR: KT-61970
2023-09-07 19:36:41 +02:00
parent 84bd12c667
commit 63e9dd588c
70 changed files with 380 additions and 153 deletions
@@ -0,0 +1,78 @@
+/*
+ * 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.*
+import org.jetbrains.kotlin.utils.SmartPrinter
+import org.jetbrains.kotlin.utils.withIndent
+
+/**
+ * A common class for printing FIR or IR tree elements.
+ */
+abstract class AbstractElementPrinter<Element : AbstractElement<Element, Field>, Field : AbstractField>(
+    private val printer: SmartPrinter,
+) {
+
+    protected abstract fun makeFieldPrinter(printer: SmartPrinter): AbstractFieldPrinter<Field>
+
+    context(ImportCollector)
+    protected abstract fun SmartPrinter.printAdditionalMethods(element: Element)
+
+    protected open fun defaultElementKDoc(element: Element): String? = null
+
+    protected open val separateFieldsWithBlankLine: Boolean
+        get() = false
+
+    context(ImportCollector)
+    fun printElement(element: Element) {
+        printer.run {
+            val kind = element.kind ?: error("Expected non-null element kind")
+
+            printKDoc(element.extendedKDoc(defaultElementKDoc(element)))
+            print(kind.title, " ", element.typeName)
+            print(element.params.typeParameters())
+
+            val parentRefs = element.parentRefs
+            if (parentRefs.isNotEmpty()) {
+                print(
+                    parentRefs.sortedBy { it.typeKind }.joinToString(prefix = " : ") { parent ->
+                        parent.render() + parent.inheritanceClauseParenthesis()
+                    }
+                )
+            }
+            print(element.params.multipleUpperBoundsList())
+
+            val body = SmartPrinter(StringBuilder()).apply {
+                val fieldPrinter = makeFieldPrinter(this)
+                withIndent {
+                    for (field in element.allFields) {
+                        if (
+                            !field.withGetter && field.defaultValueInImplementation == null && field.isFinal && field.fromParent ||
+                            field.isParameter
+                        ) {
+                            continue
+                        }
+                        if (separateFieldsWithBlankLine) println()
+                        fieldPrinter.printField(
+                            field,
+                            override = field.fromParent,
+                            modality = Modality.ABSTRACT.takeIf { !field.isFinal && !kind.isInterface },
+                        )
+                    }
+                    printAdditionalMethods(element)
+                }
+            }.toString()
+
+            if (body.isNotEmpty()) {
+                println(" {")
+                print(body)
+                print("}")
+            }
+            println()
+        }
+    }
+}
@@ -146,4 +146,120 @@ fun SmartPrinter.printFunctionDeclaration(
        print(": ", returnType.render())
    }
    print(typeParameters.multipleUpperBoundsList())
-}
+}
+
+private val dataTP = TypeVariable("D")
+private val dataParameter = FunctionParameter("data", dataTP)
+
+context(ImportCollector)
+fun SmartPrinter.printAcceptMethod(
+    element: AbstractElement<*, *>,
+    visitorClass: ClassRef<PositionTypeParameterRef>,
+    hasImplementation: Boolean,
+    kDoc: String?,
+) {
+    if (!element.hasAcceptMethod) return
+    println()
+    printKDoc(kDoc)
+    val resultTP = TypeVariable("R")
+    printFunctionDeclaration(
+        name = "accept",
+        parameters = listOf(
+            FunctionParameter("visitor", visitorClass.withArgs(resultTP, dataTP)),
+            dataParameter,
+        ),
+        returnType = resultTP,
+        typeParameters = listOf(resultTP, dataTP),
+        override = !element.isRootElement,
+    )
+    if (hasImplementation) {
+        println(" =")
+        withIndent {
+            print("visitor.", element.visitFunctionName, "(this, ", dataParameter.name, ")")
+        }
+    }
+    println()
+}
+
+context(ImportCollector)
+fun SmartPrinter.printTransformMethod(
+    element: AbstractElement<*, *>,
+    transformerClass: ClassRef<PositionTypeParameterRef>,
+    implementation: String?,
+    returnType: TypeRefWithNullability,
+    kDoc: String?,
+) {
+    if (!element.hasTransformMethod) return
+    println()
+    printKDoc(kDoc)
+    if (returnType is TypeParameterRef && implementation != null) {
+        println("@Suppress(\"UNCHECKED_CAST\")")
+    }
+    printFunctionDeclaration(
+        name = "transform",
+        parameters = listOf(
+            FunctionParameter("transformer", transformerClass.withArgs(dataTP)),
+            dataParameter,
+        ),
+        returnType = returnType,
+        typeParameters = listOfNotNull(returnType as? TypeVariable, dataTP),
+        override = !element.isRootElement,
+    )
+    if (implementation != null) {
+        println(" =")
+        withIndent {
+            print(implementation, " as ", returnType.render())
+        }
+    }
+    println()
+}
+
+context(ImportCollector)
+fun SmartPrinter.printAcceptChildrenMethod(
+    element: FieldContainer,
+    visitorClass: ClassRef<PositionTypeParameterRef>,
+    visitorResultType: TypeRef,
+    modality: Modality? = null,
+    override: Boolean = false,
+    kDoc: String?,
+) {
+    if (!element.hasAcceptChildrenMethod) return
+    println()
+    printKDoc(kDoc)
+    printFunctionDeclaration(
+        name = "acceptChildren",
+        parameters = listOf(
+            FunctionParameter("visitor", visitorClass.withArgs(visitorResultType, dataTP)),
+            dataParameter,
+        ),
+        returnType = StandardTypes.unit,
+        typeParameters = listOfNotNull(visitorResultType as? TypeVariable, dataTP),
+        modality = modality,
+        override = override,
+    )
+}
+
+context(ImportCollector)
+fun SmartPrinter.printTransformChildrenMethod(
+    element: FieldContainer,
+    transformerClass: ClassRef<PositionTypeParameterRef>,
+    returnType: TypeRef,
+    modality: Modality? = null,
+    override: Boolean = false,
+    kDoc: String?,
+) {
+    if (!element.hasTransformChildrenMethod) return
+    println()
+    printKDoc(kDoc)
+    printFunctionDeclaration(
+        name = "transformChildren",
+        parameters = listOf(
+            FunctionParameter("transformer", transformerClass.withArgs(dataTP)),
+            dataParameter,
+        ),
+        returnType = returnType,
+        typeParameters = listOf(dataTP),
+        modality = modality,
+        override = override,
+    )
+}