[FIR generator] Use TypeVariable instead of TypeArgument for TPs

This is a step towards commonizing the code generator between
FIR and IR: KT-61970

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

@@ -786,8 +786,8 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
 // TODO: Replace with org.jetbrains.kotlin.generators.tree.withArgs
 fun Element.withArgs(vararg replacements: Pair<String, TypeRef>): AbstractElement {
     val replaceMap = replacements.toMap()
-    val newArguments = typeArguments.map { typeArgument ->
+    val newArguments = params.map { param ->
-        replaceMap[typeArgument.name]?.let { SimpleTypeArgument(it.type, null) } ?: typeArgument
+        replaceMap[param.name]?.let { SimpleTypeArgument(it.type, null) } ?: error("Type variable $param not found in $this")
     }
     return ElementWithArguments(this, newArguments)
 }

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.context
 
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.types.Variance
 
 abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val builder: T) {
     inner class ConfigureContext(val element: Element) {
@@ -27,20 +28,10 @@ abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val
             }
         }
 
-        fun withArg(name: String) {
+        fun withArg(name: String, vararg upperBounds: TypeRef): TypeVariable = withArg(name, upperBounds.toList())
-            withArg(name, emptyList())
-        }
 
-        fun withArg(name: String, upperBound: TypeRef, vararg upperBounds: TypeRef) {
+        fun withArg(name: String, upperBounds: List<TypeRef>, variance: Variance = Variance.INVARIANT): TypeVariable {
-            val allUpperBounds = mutableListOf(upperBound).apply { this += upperBounds }
+            return TypeVariable(name, upperBounds, variance).also(element.params::add)
-            withArg(name, allUpperBounds)
-        }
-
-        private fun withArg(name: String, upperBounds: List<TypeRef>) {
-            element.typeArguments += when (upperBounds.size) {
-                in 0..1 -> SimpleTypeArgument(name, upperBounds.firstOrNull())
-                else -> TypeArgumentWithMultipleUpperBounds(name, upperBounds.toList())
-            }
         }
 
         fun parentArg(parent: Element, argument: String, type: TypeRef) {

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

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.generators.tree.FieldContainer
 import org.jetbrains.kotlin.generators.tree.Importable
-import org.jetbrains.kotlin.generators.tree.generics
+import org.jetbrains.kotlin.generators.tree.printer.generics
 
 private const val DEFAULT_BUILDER_PACKAGE = "org.jetbrains.kotlin.fir.tree.builder"
 

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

@@ -39,7 +39,12 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
     override val parents = mutableListOf<Element>()
     override var defaultImplementation: Implementation? = null
     override val customImplementations = mutableListOf<Implementation>()
-    override val typeArguments = mutableListOf<TypeArgument>()
+
+    override val params = mutableListOf<TypeVariable>()
+
+    override val typeArguments: List<TypeArgument>
+        get() = emptyList()
+
     override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<TypeRef, TypeRef>>()
     override var kind: ImplementationKind? = null
         set(value) {

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.fir.tree.generator.model
 
 import org.jetbrains.kotlin.fir.tree.generator.printer.generics
 import org.jetbrains.kotlin.generators.tree.*
+import org.jetbrains.kotlin.generators.tree.printer.generics
 
 class ImplementationWithArg(
     val implementation: Implementation,

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

@@ -248,8 +248,8 @@ private fun SmartPrinter.printDslBuildFunction(
         println("@OptIn(FirImplementationDetail::class)")
     }
     print("fun ")
-    builder.implementation.element.typeArguments.takeIf { it.isNotEmpty() }?.let {
+    builder.implementation.element.params.takeIf { it.isNotEmpty() }?.let {
-        print(it.joinToString(separator = ", ", prefix = "<", postfix = "> ") { it.name })
+        print(it.joinToString(separator = ", ", prefix = "<", postfix = "> "))
     }
     val builderType = builder.typeWithArguments
     val name = builder.implementation.name?.replaceFirst("Fir", "") ?: builder.implementation.element.name
@@ -288,8 +288,8 @@ private fun SmartPrinter.printDslBuildCopyFunction(
     println("@OptIn(${optIns.joinToString { "$it::class" }})")
     print("inline ")
     print("fun ")
-    builder.implementation.element.typeArguments.takeIf { it.isNotEmpty() }?.let {
+    builder.implementation.element.params.takeIf { it.isNotEmpty() }?.let {
-        print(it.joinToString(separator = ", ", prefix = "<", postfix = "> ") { it.name })
+        print(it.joinToString(separator = ", ", prefix = "<", postfix = "> "))
     }
     val builderType = builder.typeWithArguments
     val name = builder.implementation.name?.replaceFirst("Fir", "") ?: builder.implementation.element.name

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

@@ -8,12 +8,12 @@ package org.jetbrains.kotlin.fir.tree.generator.printer
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.model.Element
 import org.jetbrains.kotlin.fir.tree.generator.model.Field
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
-import org.jetbrains.kotlin.generators.tree.Importable
 import org.jetbrains.kotlin.fir.tree.generator.pureAbstractElementType
 import org.jetbrains.kotlin.fir.tree.generator.util.get
-import org.jetbrains.kotlin.generators.tree.TypeRef
+import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.generators.tree.typeWithArguments
+import org.jetbrains.kotlin.generators.tree.printer.braces
+import org.jetbrains.kotlin.generators.tree.printer.multipleUpperBoundsList
+import org.jetbrains.kotlin.generators.tree.printer.typeParameters
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
 import java.io.File
@@ -53,9 +53,7 @@ fun SmartPrinter.printElement(element: Element) {
         }
 
         print("${kind!!.title} $type")
-        if (typeArguments.isNotEmpty()) {
+        print(typeParameters())
-            print(typeArguments.joinToString(", ", "<", ">") { it.toString() })
-        }
         val needPureAbstractElement = !isInterface && !allParents.any { it.kind == ImplementationKind.AbstractClass || it.kind == ImplementationKind.SealedClass }
 
         if (parents.isNotEmpty() || needPureAbstractElement) {

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

@@ -6,11 +6,10 @@
 package org.jetbrains.kotlin.fir.tree.generator.printer
 
 import org.jetbrains.kotlin.fir.tree.generator.model.*
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
 import org.jetbrains.kotlin.fir.tree.generator.pureAbstractElementType
-import org.jetbrains.kotlin.generators.tree.Importable
+import org.jetbrains.kotlin.generators.tree.*
-import org.jetbrains.kotlin.generators.tree.TypeRef
+import org.jetbrains.kotlin.generators.tree.printer.braces
-import org.jetbrains.kotlin.generators.tree.typeWithArguments
+import org.jetbrains.kotlin.generators.tree.printer.typeParameters
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
 import org.jetbrains.kotlin.utils.withIndent
@@ -70,7 +69,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
             print("internal ")
         }
         print("${kind!!.title} $type")
-        print(element.typeParameters)
+        print(element.typeParameters(end = " "))
 
         val isInterface = kind == ImplementationKind.Interface || kind == ImplementationKind.SealedInterface
         val isAbstract = kind == ImplementationKind.AbstractClass || kind == ImplementationKind.SealedClass

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

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.fir.tree.generator.printer
 
 import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.model.Element
+import org.jetbrains.kotlin.generators.tree.printer.multipleUpperBoundsList
+import org.jetbrains.kotlin.generators.tree.printer.typeParameters
 import org.jetbrains.kotlin.generators.tree.typeWithArguments
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
@@ -34,7 +36,7 @@ fun printTransformer(elements: List<Element>, generationPath: File): GeneratedFi
                 if (element == AbstractFirTreeBuilder.baseFirElement) continue
                 val varName = element.safeDecapitalizedName
                 print("open fun ")
-                element.typeParameters.takeIf { it.isNotBlank() }?.let { print(it) }
+                element.typeParameters(end = " ").takeIf { it.isNotBlank() }?.let { print(it) }
                 println(
                     "transform${element.name}($varName: ${element.typeWithArguments}, data: D): ${element.transformerType
                         .typeWithArguments}${element.multipleUpperBoundsList()} {",
@@ -49,7 +51,7 @@ fun printTransformer(elements: List<Element>, generationPath: File): GeneratedFi
             for (element in elements) {
                 val varName = element.safeDecapitalizedName
                 print("final override fun ")
-                element.typeParameters.takeIf { it.isNotBlank() }?.let { print(it) }
+                element.typeParameters(end = " ").takeIf { it.isNotBlank() }?.let { print(it) }
 
                 println(
                     "visit${element.name}($varName: ${element.typeWithArguments}, data: D): ${element.transformerType

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

@@ -77,7 +77,7 @@ private fun Element.collectImportsInternal(base: List<String>, kind: ImportKind)
     val fqns = base + allFields.mapNotNull { it.fullQualifiedName } +
             allFields.flatMap { it.overridenTypes.mapNotNull { it.fullQualifiedName } + it.arbitraryImportables.mapNotNull { it.fullQualifiedName } } +
             allFields.flatMap { it.arguments.mapNotNull { it.fullQualifiedName } } +
-            typeArguments.flatMap { it.upperBounds.mapNotNull { it.fullQualifiedName } }
+            params.flatMap { it.bounds.mapNotNull { it.fullQualifiedName } }
     val result = fqns.filterRedundantImports(packageName, kind).toMutableList()
 
     if (allFields.any { it is FieldList && it.isMutableOrEmpty }) {
@@ -149,21 +149,6 @@ fun Field.getMutableType(forBuilder: Boolean = false, notNull: Boolean = false):
 
 fun Field.call(): String = if (nullable) "?." else "."
 
-fun Element.multipleUpperBoundsList(): String {
-    return typeArguments.filterIsInstance<TypeArgumentWithMultipleUpperBounds>().takeIf { it.isNotEmpty() }?.let { arguments ->
-        val upperBoundsList = arguments.joinToString(", ") { argument ->
-            argument.upperBounds.joinToString(", ") { upperBound -> "${argument.name} : ${upperBound.typeWithArguments}" }
-        }
-        " where $upperBoundsList"
-    } ?: ""
-}
-
-fun ImplementationKind?.braces(): String = when (this) {
-    ImplementationKind.Interface, ImplementationKind.SealedInterface -> ""
-    ImplementationKind.OpenClass, ImplementationKind.AbstractClass, ImplementationKind.SealedClass -> "()"
-    else -> throw IllegalStateException(this.toString())
-}
-
 val Element.safeDecapitalizedName: String get() = if (name == "Class") "klass" else name.replaceFirstChar(Char::lowercaseChar)
 
 val ImplementationWithArg.generics: String
@@ -173,8 +158,3 @@ val Field.generics: String
     get() = arguments.takeIf { it.isNotEmpty() }
         ?.let { it.joinToString(", ", "<", ">") { it.typeWithArguments } }
         ?: ""
-
-val Element.typeParameters: String
-    get() = typeArguments.takeIf { it.isNotEmpty() }
-        ?.joinToString(", ", "<", "> ")
-        ?: ""

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

@@ -8,6 +8,8 @@ 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.model.Element
+import org.jetbrains.kotlin.generators.tree.printer.multipleUpperBoundsList
+import org.jetbrains.kotlin.generators.tree.printer.typeParameters
 import org.jetbrains.kotlin.generators.tree.typeWithArguments
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
@@ -61,7 +63,7 @@ fun printVisitor(elements: List<Element>, generationPath: File, visitSuperTypeBy
                 } else {
                     print("open")
                 }
-                print(" fun ${typeParameters}visit$name($varName: $typeWithArguments, data: D): R${multipleUpperBoundsList()} = visit")
+                print(" fun ${typeParameters(end = " ")}visit$name($varName: $typeWithArguments, data: D): R${multipleUpperBoundsList()} = visit")
                 if (visitSuperTypeByDefault) {
                     print(element.getNameOfSupertypeForDefaultVisiting())
                 } else {
@@ -99,7 +101,7 @@ fun printVisitorVoid(elements: List<Element>, generationPath: File): GeneratedFi
                 if (element == AbstractFirTreeBuilder.baseFirElement) continue
                 with(element) {
                     val varName = safeDecapitalizedName
-                    println("open fun ${typeParameters}visit$name($varName: $typeWithArguments)${multipleUpperBoundsList()} {")
+                    println("open fun ${typeParameters(end = " ")}visit$name($varName: $typeWithArguments)${multipleUpperBoundsList()} {")
                     withIndent {
                         println("visitElement($varName)")
                     }
@@ -111,7 +113,7 @@ fun printVisitorVoid(elements: List<Element>, generationPath: File): GeneratedFi
             for (element in elements) {
                 with(element) {
                     val varName = safeDecapitalizedName
-                    println("final override fun ${typeParameters}visit$name($varName: $typeWithArguments, data: Nothing?)${multipleUpperBoundsList()} {")
+                    println("final override fun ${typeParameters(end = " ")}visit$name($varName: $typeWithArguments, data: Nothing?)${multipleUpperBoundsList()} {")
                     withIndent {
                         println("visit$name($varName)")
                     }
@@ -145,7 +147,7 @@ fun printDefaultVisitorVoid(elements: List<Element>, generationPath: File): Gene
             if (!element.isAcceptableForDefaultVisiting()) continue
             with(element) {
                 val varName = safeDecapitalizedName
-                println("override fun ${typeParameters}visit$name($varName: $typeWithArguments)${multipleUpperBoundsList()} = visit${element.getNameOfSupertypeForDefaultVisiting()}($varName)")
+                println("override fun ${typeParameters(end = " ")}visit$name($varName: $typeWithArguments)${multipleUpperBoundsList()} = visit${element.getNameOfSupertypeForDefaultVisiting()}($varName)")
                 println()
             }
         }

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.generators.tree
 
+import org.jetbrains.kotlin.generators.tree.printer.generics
+
 /**
  * A common interface representing a FIR or IR tree element.
  */
@@ -17,6 +19,10 @@ interface AbstractElement<Element : AbstractElement<Element, Field>, Field : Abs
 
     val parents: List<Element>
 
+    val params: List<TypeVariable>
+
+    // TODO: Remove this property as soon as we replace org.jetbrains.kotlin.fir.tree.generator.model.ElementWithArguments with
+    // a generic ElementRef class
     val typeArguments: List<TypeArgument>
 
     val parentsArguments: Map<Element, Map<TypeRef, TypeRef>>
@@ -31,8 +37,3 @@ interface AbstractElement<Element : AbstractElement<Element, Field>, Field : Abs
 
     override fun getTypeWithArguments(notNull: Boolean): String = type + generics
 }
-
-val AbstractElement<*, *>.generics: String
-    get() = typeArguments.takeIf { it.isNotEmpty() }
-        ?.let { it.joinToString(", ", "<", ">") { it.name } }
-        ?: ""

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

@@ -0,0 +1,54 @@
+/*
+ * 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.printer
+
+import org.jetbrains.kotlin.generators.tree.AbstractElement
+import org.jetbrains.kotlin.generators.tree.ImplementationKind
+import org.jetbrains.kotlin.generators.tree.typeWithArguments
+
+/**
+ * The angle bracket-delimited list of type parameters to print, or empty string if the element has no type parameters.
+ *
+ * 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]).
+ *
+ * @param end The string to add after the closing angle bracket of the type parameter list
+ */
+fun AbstractElement<*, *>.typeParameters(end: String = ""): String = params.takeIf { it.isNotEmpty() }
+    ?.joinToString(", ", "<", ">$end") { param ->
+        param.name + (param.bounds.singleOrNull()?.let { " : ${it.typeWithArguments}" } ?: "")
+    } ?: ""
+
+/**
+ * The `where` clause to print after the class or function declaration if at least one of the element's tye parameters has multiple upper
+ * bounds.
+ *
+ * Otherwise, an empty string.
+ */
+fun AbstractElement<*, *>.multipleUpperBoundsList(): String {
+    val paramsWithMultipleUpperBounds = params.filter { it.bounds.size > 1 }.takeIf { it.isNotEmpty() } ?: return ""
+    return buildString {
+        append(" where ")
+        paramsWithMultipleUpperBounds.joinTo(this, separator = ", ") { param ->
+            param.bounds.joinToString(", ") { bound -> "$param : ${bound.typeWithArguments}" }
+        }
+        append("")
+    }
+}
+
+/**
+ * The braces to print in the inheritance clause if this is a class, or empty string if this is an interface.
+ */
+fun ImplementationKind?.braces(): String = when (this) {
+    ImplementationKind.Interface, ImplementationKind.SealedInterface -> ""
+    ImplementationKind.OpenClass, ImplementationKind.AbstractClass, ImplementationKind.SealedClass -> "()"
+    else -> throw IllegalStateException(this.toString())
+}
+
+val AbstractElement<*, *>.generics: String
+    get() = params.takeIf { it.isNotEmpty() }
+        ?.let { it.joinToString(", ", "<", ">") { it.name } }
+        ?: ""