[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
@@ -38,6 +38,21 @@ abstract class AbstractElement<Element, Field> : ElementOrRef<Element, Field>, F
    open val isSealed: Boolean
        get() = false

+    /**
+     * The name of the method in visitors used to visit this element.
+     */
+    abstract val visitFunctionName: String
+
+    /**
+     * The name of the parameter representing this element in the visitor method used to visit this element.
+     */
+    abstract val visitorParameterName: String
+
+    /**
+     * The default element to visit if the method for visiting this element is not overridden.
+     */
+    abstract val parentInVisitor: Element?
+
    override val allParents: List<Element>
        get() = elementParents.map { it.element }

@@ -0,0 +1,159 @@
+/*
+ * 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.multipleUpperBoundsList
+import org.jetbrains.kotlin.generators.tree.printer.printFunctionDeclaration
+import org.jetbrains.kotlin.generators.tree.printer.typeParameters
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.SmartPrinter
+import org.jetbrains.kotlin.utils.withIndent
+
+abstract class AbstractVisitorPrinter<Element : AbstractElement<Element, Field>, Field : AbstractField>(
+    val printer: SmartPrinter,
+    val visitSuperTypeByDefault: Boolean,
+) {
+
+    /**
+     * The visitor type to print.
+     */
+    abstract val visitorType: ClassRef<*>
+
+    /**
+     * The result type parameter of the visitor. All visitor methods return result of this type.
+     */
+    protected val resultTypeVariable = TypeVariable("R", emptyList(), Variance.OUT_VARIANCE)
+
+    /**
+     * The data type parameter of the visitor. ALl visitor methods accept a parameter of this type.
+     */
+    protected val dataTypeVariable = TypeVariable("D", emptyList(), Variance.IN_VARIANCE)
+
+    /**
+     * The type parameters of the visitor class. Void visitors have no type parameters,
+     * regular visitors usually have [resultTypeVariable] and [dataTypeVariable] here.
+     */
+    abstract val visitorTypeParameters: List<TypeVariable>
+
+    abstract val visitorDataType: TypeRef
+
+    abstract fun visitMethodReturnType(element: Element): TypeRef
+
+    /**
+     * The superclass for this visitor class.
+     */
+    abstract val visitorSuperType: ClassRef<PositionTypeParameterRef>?
+
+    /**
+     * If `true`, visitor methods for generic tree elements will be parameterized correspondingly.
+     * Otherwise, type arguments of generic tree elements will be replaced with `*`.
+     */
+    abstract val allowTypeParametersInVisitorMethods: Boolean
+
+    /**
+     * Allows to customize the default element to visit if the method for visiting this [element] is not overridden.
+     *
+     * If returns `null`, methods for this element will not be overridden in this visitor class (except the root element).
+     */
+    open fun parentInVisitor(element: Element): Element? = element.parentInVisitor
+
+    /**
+     * Prints a single visitor method declaration, without body.
+     */
+    context(ImportCollector)
+    protected fun SmartPrinter.printVisitMethodDeclaration(
+        element: Element,
+        hasDataParameter: Boolean = true,
+        modality: Modality? = null,
+        override: Boolean = false,
+    ) {
+        val visitorParameterType = ElementRef(
+            element,
+            element.params.associateWith { if (allowTypeParametersInVisitorMethods) it else TypeRef.Star }
+        )
+        val parameters = buildList {
+            add(FunctionParameter(element.visitorParameterName, visitorParameterType))
+            if (hasDataParameter) add(FunctionParameter("data", visitorDataType))
+        }
+        printFunctionDeclaration(
+            name = element.visitFunctionName,
+            parameters = parameters,
+            returnType = visitMethodReturnType(element),
+            typeParameters = if (allowTypeParametersInVisitorMethods) {
+                element.params
+            } else {
+                emptyList()
+            },
+            modality = modality,
+            override = override,
+        )
+    }
+
+    context(ImportCollector)
+    protected fun printMethodDeclarationForElement(element: Element, modality: Modality? = null, override: Boolean) {
+        printer.run {
+            println()
+            printVisitMethodDeclaration(
+                element,
+                modality = modality,
+                override = override
+            )
+        }
+    }
+
+    context(ImportCollector)
+    protected open fun printMethodsForElement(element: Element) {
+        printer.run {
+            val parentInVisitor = parentInVisitor(element)
+            if (parentInVisitor == null && !element.isRootElement) return
+            printMethodDeclarationForElement(
+                element,
+                modality = when {
+                    visitorSuperType == null && parentInVisitor == null && visitorType.kind == TypeKind.Class -> Modality.ABSTRACT
+                    visitorSuperType == null && parentInVisitor != null && visitorType.kind == TypeKind.Class -> Modality.OPEN
+                    else -> null
+                },
+                override = parentInVisitor != null && visitorSuperType != null,
+            )
+            if (parentInVisitor != null) {
+                print(" = ", parentInVisitor.visitFunctionName, "(", element.visitorParameterName, ", data)")
+            }
+            println()
+        }
+    }
+
+    context(ImportCollector)
+    protected open fun SmartPrinter.printAdditionalMethods() {
+    }
+
+    context(ImportCollector)
+    fun printVisitor(elements: List<Element>) {
+        val visitorType = this.visitorType
+        printer.run {
+            when (visitorType.kind) {
+                TypeKind.Interface -> print("interface ")
+                TypeKind.Class -> print("abstract class ")
+            }
+            print(visitorType.simpleName, visitorTypeParameters.typeParameters())
+            visitorSuperType?.let {
+                print(" : ", it.render(), it.inheritanceClauseParenthesis())
+            }
+            print(visitorTypeParameters.multipleUpperBoundsList())
+            println(" {")
+            withIndent {
+                printAdditionalMethods()
+                for (element in elements) {
+                    if (element.isRootElement && visitSuperTypeByDefault) continue
+                    if (visitSuperTypeByDefault && parentInVisitor(element) == null) continue
+                    printMethodsForElement(element)
+                }
+            }
+            println("}")
+        }
+    }
+}
@@ -0,0 +1,81 @@
+/*
+ * 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.utils.SmartPrinter
+import org.jetbrains.kotlin.utils.withIndent
+
+abstract class AbstractVisitorVoidPrinter<Element, Field>(
+    printer: SmartPrinter,
+    visitSuperTypeByDefault: Boolean,
+) : AbstractVisitorPrinter<Element, Field>(printer, visitSuperTypeByDefault)
+        where Element : AbstractElement<Element, Field>,
+              Field : AbstractField {
+
+    final override val visitorTypeParameters: List<TypeVariable>
+        get() = emptyList()
+
+    final override val visitorDataType: TypeRef
+        get() = StandardTypes.nothing.copy(nullable = true)
+
+    override fun visitMethodReturnType(element: Element) = StandardTypes.unit
+
+    abstract val visitorSuperClass: ClassRef<PositionTypeParameterRef>
+
+    override val visitorSuperType: ClassRef<PositionTypeParameterRef>
+        get() = if (visitSuperTypeByDefault)
+            visitorSuperClass
+        else
+            visitorSuperClass.withArgs(StandardTypes.unit, visitorDataType)
+
+    abstract val useAbstractMethodForRootElement: Boolean
+
+    abstract val overriddenVisitMethodsAreFinal: Boolean
+
+    context(ImportCollector)
+    final override fun printMethodsForElement(element: Element) {
+        val parentInVisitor = parentInVisitor(element)
+        if (!element.isRootElement && parentInVisitor == null) return
+
+        val isAbstractVisitRootElementMethod = element.isRootElement && useAbstractMethodForRootElement
+
+        printMethodDeclarationForElement(
+            element,
+            modality = Modality.FINAL.takeIf { overriddenVisitMethodsAreFinal },
+            override = true,
+        )
+
+        fun SmartPrinter.printBody(parentInVisitor: Element?) {
+            println(" {")
+            if (parentInVisitor != null) {
+                withIndent {
+                    println(parentInVisitor.visitFunctionName, "(", element.visitorParameterName, ")")
+                }
+            }
+            println("}")
+        }
+
+        printer.run {
+            printBody(element)
+            println()
+            printVisitMethodDeclaration(
+                element,
+                hasDataParameter = false,
+                modality = when {
+                    element.isRootElement && visitorType.kind == TypeKind.Class -> Modality.ABSTRACT
+                    !element.isRootElement && visitorType.kind == TypeKind.Class -> Modality.OPEN
+                    else -> null
+                }
+            )
+            if (isAbstractVisitRootElementMethod) {
+                println()
+            } else {
+                printBody(parentInVisitor)
+            }
+        }
+    }
+}
@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.generators.tree

 object StandardTypes {
+    val unit = type<Unit>()
+    val nothing = type("kotlin", "Nothing")
    val boolean = type<Boolean>()
    val string = type<String>()
    val int = type<Int>()
@@ -5,14 +5,17 @@

 package org.jetbrains.kotlin.generators.tree.printer

-import org.jetbrains.kotlin.generators.tree.AbstractElement
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.generators.tree.ImportCollector
 import org.jetbrains.kotlin.generators.tree.render
-import org.jetbrains.kotlin.utils.SmartPrinter
+import org.jetbrains.kotlin.types.Variance
+import org.jetbrains.kotlin.utils.addToStdlib.joinToWithBuffer
+import org.jetbrains.kotlin.utils.withIndent

 /**
- * The angle bracket-delimited list of type parameters to print, or empty string if the element has no type parameters.
+ * The angle bracket-delimited list of type parameters to print, or empty string if the list is empty.
 *
 * For type parameters that have a single upper bound, also prints that upper bound. If at least one type parameter has multiple upper
 * bounds, doesn't print any upper bounds at all. They are expected to be printed in the `where` clause (see [multipleUpperBoundsList]).
@@ -20,10 +23,21 @@ import org.jetbrains.kotlin.utils.SmartPrinter
 * @param end The string to add after the closing angle bracket of the type parameter list
 */
 context(ImportCollector)
-fun AbstractElement<*, *>.typeParameters(end: String = ""): String = params.takeIf { it.isNotEmpty() }
-    ?.joinToString(", ", "<", ">$end") { param ->
-        param.name + (param.bounds.singleOrNull()?.let { " : ${it.render()}" } ?: "")
-    } ?: ""
+fun List<TypeVariable>.typeParameters(end: String = ""): String = buildString {
+    if (this@typeParameters.isEmpty()) return@buildString
+    joinToWithBuffer(this, prefix = "<", postfix = ">") { param ->
+        if (param.variance != Variance.INVARIANT) {
+            append(param.variance.label)
+            append(" ")
+        }
+        append(param.name)
+        param.bounds.singleOrNull()?.let {
+            append(" : ")
+            it.renderTo(this)
+        }
+    }
+    append(end)
+}

 /**
 * The `where` clause to print after the class or function declaration if at least one of the element's tye parameters has multiple upper
@@ -32,14 +46,17 @@ fun AbstractElement<*, *>.typeParameters(end: String = ""): String = params.take
 * Otherwise, an empty string.
 */
 context(ImportCollector)
-fun AbstractElement<*, *>.multipleUpperBoundsList(): String {
-    val paramsWithMultipleUpperBounds = params.filter { it.bounds.size > 1 }.takeIf { it.isNotEmpty() } ?: return ""
+fun List<TypeVariable>.multipleUpperBoundsList(): String {
+    val paramsWithMultipleUpperBounds = filter { it.bounds.size > 1 }.takeIf { it.isNotEmpty() } ?: return ""
    return buildString {
        append(" where ")
-        paramsWithMultipleUpperBounds.joinTo(this, separator = ", ") { param ->
-            param.bounds.joinToString(", ") { bound -> "$param : ${bound.render()}" }
+        paramsWithMultipleUpperBounds.joinToWithBuffer(this, separator = ", ") { param ->
+            param.bounds.joinToWithBuffer(this) { bound ->
+                append(param.name)
+                append(" : ")
+                bound.renderTo(this)
+            }
        }
-        append("")
    }
 }

@@ -75,3 +92,58 @@ fun AbstractElement<*, *>.extendedKDoc(defaultKDoc: String? = null): String = bu
    }
    append("Generated from: [${element.propertyName}]")
 }
+
+data class FunctionParameter(val name: String, val type: TypeRef, val defaultValue: String? = null) {
+
+    context(ImportCollector)
+    fun render(): String = buildString {
+        append(name, ": ", type.render())
+        defaultValue?.let {
+            append(" = ", it)
+        }
+    }
+}
+
+context(ImportCollector)
+fun SmartPrinter.printFunctionDeclaration(
+    name: String,
+    parameters: List<FunctionParameter>,
+    returnType: TypeRef,
+    typeParameters: List<TypeVariable> = emptyList(),
+    modality: Modality? = null,
+    override: Boolean = false,
+    allParametersOnSeparateLines: Boolean = false,
+) {
+    when (modality) {
+        null -> {}
+        Modality.FINAL -> print("final ")
+        Modality.OPEN -> print("open ")
+        Modality.ABSTRACT -> print("abstract ")
+        Modality.SEALED -> error("Function cannot be sealed")
+    }
+    if (override) {
+        print("override ")
+    }
+    print("fun ")
+    print(typeParameters.typeParameters(end = " "))
+    print(name, "(")
+
+    if (allParametersOnSeparateLines) {
+        if (parameters.isNotEmpty()) {
+            println()
+            withIndent {
+                for (parameter in parameters) {
+                    print(parameter.render())
+                    println(",")
+                }
+            }
+        }
+    } else {
+        print(parameters.joinToString { it.render() })
+    }
+    print(")")
+    if (returnType != StandardTypes.unit) {
+        print(": ", returnType.render())
+    }
+    print(typeParameters.multipleUpperBoundsList())
+}