FIR tree: support const expressions (adds IR dependency) #KT-24023 Fixed

Also support generic tree elements in FIR visitor generator

compiler/fir/psi2fir/build.gradle.kts

@@ -14,6 +14,7 @@ dependencies {
     compile(project(":compiler:psi"))
     compile(project(":core:descriptors"))
     compile(project(":compiler:fir:tree"))
+    compile(project(":compiler:ir.tree"))
 
     compileOnly(intellijCoreDep()) { includeJars("intellij-core", "annotations") }
     

compiler/fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.builder
 
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -24,6 +25,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
 import org.jetbrains.kotlin.fir.types.FirType
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.impl.*
+import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.name.ClassId
@@ -91,8 +93,127 @@ class RawFirBuilder(val session: FirSession) {
                 else -> FirExpressionBodyImpl(session, FirExpressionStub(session, null))
             }
 
-        private fun ValueArgument?.toFirExpression(): FirExpression = with(this as? KtElement) {
+        private fun String.parseCharacter(): Char? {
-            convertSafe() ?: FirErrorExpressionImpl(session, this)
+            // Strip the quotes
+            if (length < 2 || this[0] != '\'' || this[length - 1] != '\'') {
+                return null
+            }
+            val text = substring(1, length - 1) // now there're no quotes
+
+            if (text.isEmpty()) {
+                return null
+            }
+
+            return if (text[0] != '\\') {
+                // No escape
+                if (text.length == 1) {
+                    text[0]
+                } else {
+                    null
+                }
+            } else {
+                escapedStringToCharacter(text)
+            }
+        }
+
+        private fun escapedStringToCharacter(text: String): Char? {
+            assert(text.isNotEmpty() && text[0] == '\\') {
+                "Only escaped sequences must be passed to this routine: $text"
+            }
+
+            // Escape
+            val escape = text.substring(1) // strip the slash
+            when (escape.length) {
+                0 -> {
+                    // bare slash
+                    return null
+                }
+                1 -> {
+                    // one-char escape
+                    return translateEscape(escape[0]) ?: return null
+                }
+                5 -> {
+                    // unicode escape
+                    if (escape[0] == 'u') {
+                        try {
+                            val intValue = Integer.valueOf(escape.substring(1), 16)
+                            return intValue.toInt().toChar()
+                        } catch (e: NumberFormatException) {
+                            // Will be reported below
+                        }
+                    }
+                }
+            }
+            return null
+        }
+
+        private fun translateEscape(c: Char): Char? =
+            when (c) {
+                't' -> '\t'
+                'b' -> '\b'
+                'n' -> '\n'
+                'r' -> '\r'
+                '\'' -> '\''
+                '\"' -> '\"'
+                '\\' -> '\\'
+                '$' -> '$'
+                else -> null
+            }
+
+        private fun ValueArgument?.toFirExpression(): FirExpression {
+            this ?: return FirErrorExpressionImpl(session, this as? KtElement, "No argument given")
+            val expression = this.getArgumentExpression()
+            when (expression) {
+                is KtConstantExpression -> {
+                    val type = expression.node.elementType
+                    val text: String = expression.text
+                    return when (type) {
+                        KtNodeTypes.INTEGER_CONSTANT ->
+                            if (text.last() == 'l' || text.last() == 'L') {
+                                FirConstExpressionImpl(
+                                    session, expression, IrConstKind.Long, text.dropLast(1).toLongOrNull(), "Incorrect long: $text"
+                                )
+                            } else {
+                                // TODO: support byte / short
+                                FirConstExpressionImpl(session, expression, IrConstKind.Int, text.toIntOrNull(), "Incorrect int: $text")
+                            }
+                        KtNodeTypes.FLOAT_CONSTANT ->
+                            if (text.last() == 'f' || text.last() == 'F') {
+                                FirConstExpressionImpl(
+                                    session, expression, IrConstKind.Float, text.dropLast(1).toFloatOrNull(), "Incorrect float: $text"
+                                )
+                            } else {
+                                FirConstExpressionImpl(
+                                    session, expression, IrConstKind.Double, text.toDoubleOrNull(), "Incorrect double: $text"
+                                )
+                            }
+                        KtNodeTypes.CHARACTER_CONSTANT ->
+                            FirConstExpressionImpl(
+                                session, expression, IrConstKind.Char, text.parseCharacter(), "Incorrect character: $text"
+                            )
+                        KtNodeTypes.BOOLEAN_CONSTANT ->
+                            FirConstExpressionImpl(session, expression, IrConstKind.Boolean, text.toBoolean())
+                        KtNodeTypes.NULL ->
+                            FirConstExpressionImpl(session, expression, IrConstKind.Null, null)
+                        else ->
+                            throw AssertionError("Unknown literal type: $type, $text")
+                    }
+                }
+
+                is KtStringTemplateExpression -> {
+                    val sb = StringBuilder()
+                    for (entry in expression.entries) {
+                        when (entry) {
+                            is KtLiteralStringTemplateEntry -> sb.append(entry.text)
+                            is KtEscapeStringTemplateEntry -> sb.append(entry.unescapedValue)
+                            else -> return FirErrorExpressionImpl(session, expression, "Incorrect template entry: ${entry.text}")
+                        }
+                    }
+                    return FirConstExpressionImpl(session, expression, IrConstKind.String, sb.toString())
+                }
+
+                else -> return FirExpressionStub(session, this as? KtElement)
+            }
         }
 
         private fun KtPropertyAccessor?.toFirPropertyAccessor(

compiler/fir/resolve/testData/resolve/multifile/Annotations.1.kt

@@ -8,3 +8,5 @@ annotation class WithInt(val value: Int)
 annotation class WithString(val s: String)
 
 annotation class Complex(val wi: WithInt, val ws: WithString)
+
+annotation class VeryComplex(val f: Float, val d: Double, val b: Boolean, val l: Long, val n: Int?)

compiler/fir/resolve/testData/resolve/multifile/Annotations.kt

@@ -25,4 +25,5 @@ class Second(val y: Char) : @WithInt(0) First() {
 }
 
 @WithInt(24)
+@VeryComplex(3.14f, 6.67e-11, false, 123456789012345L, null)
 typealias Third = @Simple Second

compiler/fir/resolve/testData/resolve/multifile/Annotations.txt

@@ -1,16 +1,16 @@
 FILE: Annotations.kt
     @FILE:R|annotations/Simple|()
-    @R|annotations/WithInt|(STUB) public abstract class First {
+    @R|annotations/WithInt|(Int(42)) public abstract class First {
         public constructor(): super<R|kotlin/Any|>()
 
-        @R|annotations/Simple|() public abstract function foo(@R|annotations/WithString|(STUB) arg: @R|annotations/Simple|() R|kotlin/Double|): R|kotlin/Unit|
+        @R|annotations/Simple|() public abstract function foo(@R|annotations/WithString|(String(abc)) arg: @R|annotations/Simple|() R|kotlin/Double|): R|kotlin/Unit|
 
         @R|annotations/Complex|(STUB, STUB) public abstract property v(val): R|kotlin/String|
             public get(): R|kotlin/String|
 
     }
-    @R|annotations/WithString|(STUB) public final class Second : @R|annotations/WithInt|(STUB) R|test/First| {
+    @R|annotations/WithString|(String(xyz)) public final class Second : @R|annotations/WithInt|(Int(0)) R|test/First| {
-        public constructor(y: R|kotlin/Char|): super<@R|annotations/WithInt|(STUB) R|test/First|>()
+        public constructor(y: R|kotlin/Char|): super<@R|annotations/WithInt|(Int(0)) R|test/First|>()
 
         public final property y(val): R|kotlin/Char|
             public get(): R|kotlin/Char|
@@ -23,7 +23,7 @@ FILE: Annotations.kt
                 STUB
             }
 
-        @R|annotations/WithString|(STUB) public constructor(): this<R|test/Second|>()
+        @R|annotations/WithString|(String(constructor)) public constructor(): this<R|test/Second|>()
 
     }
-    @R|annotations/WithInt|(STUB) @R|annotations/WithInt|(STUB) public final typealias Third = @R|annotations/Simple|() R|test/Second|
+    @R|annotations/WithInt|(Int(24)) @R|annotations/VeryComplex|(Float(3.14), Double(6.67E-11), Boolean(false), Long(123456789012345), Null(null)) @R|annotations/WithInt|(Int(24)) @R|annotations/VeryComplex|(Float(3.14), Double(6.67E-11), Boolean(false), Long(123456789012345), Null(null)) public final typealias Third = @R|annotations/Simple|() R|test/Second|

compiler/fir/tree/build.gradle.kts

@@ -13,6 +13,7 @@ jvmTarget = "1.6"
 dependencies {
     compile(project(":core:descriptors"))
     compile(project(":compiler:fir:cones"))
+    compile(project(":compiler:ir.tree"))
     // Necessary only to store bound PsiElement inside FirElement
     compileOnly(intellijCoreDep()) { includeJars("intellij-core", "annotations") }
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirRenderer.kt

@@ -361,6 +361,10 @@ class FirRenderer(builder: StringBuilder) : FirVisitorVoid() {
         print("STUB")
     }
 
+    override fun <T> visitConstExpression(constExpression: FirConstExpression<T>) {
+        print("${constExpression.kind}(${constExpression.value})")
+    }
+
     override fun visitCall(call: FirCall) {
         print("(")
         call.arguments.renderSeparated()

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirConstExpression.kt

@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.fir.expressions
+
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.ir.expressions.IrConstKind
+
+interface FirConstExpression<T> : FirExpression {
+    val kind: IrConstKind<T>
+    val value: T
+
+    override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R {
+        return visitor.visitConstExpression(this, data)
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirErrorExpression.kt

@@ -5,4 +5,6 @@
 
 package org.jetbrains.kotlin.fir.expressions
 
-interface FirErrorExpression : FirExpression
+interface FirErrorExpression : FirExpression {
+    val reason: String
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirConstExpressionImpl.kt

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.fir.expressions.impl
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.expressions.FirConstExpression
+import org.jetbrains.kotlin.ir.expressions.IrConstKind
+
+class FirConstExpressionImpl<T>(
+    override val session: FirSession,
+    override val psi: PsiElement?,
+    override val kind: IrConstKind<T>,
+    override val value: T
+) : FirConstExpression<T>
+
+fun <T> FirConstExpressionImpl(session: FirSession, psi: PsiElement?, kind: IrConstKind<T>, value: T?, errorReason: String) =
+    value?.let { FirConstExpressionImpl(session, psi, kind, it) } ?: FirErrorExpressionImpl(session, psi, errorReason)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirErrorExpressionImpl.kt

@@ -11,5 +11,6 @@ import org.jetbrains.kotlin.fir.expressions.FirErrorExpression
 
 class FirErrorExpressionImpl(
     override val session: FirSession,
-    override val psi: PsiElement?
+    override val psi: PsiElement?,
+    override val reason: String
 ) : FirErrorExpression

compiler/fir/tree/visitors-generator/src/AbstractVisitorGenerator.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.utils.Printer
 abstract class AbstractVisitorGenerator(val referencesData: DataCollector.ReferencesData) {
     fun Printer.generateFunction(
         name: String,
-        parameters: Map<String, String>,
+        parameters: Map<String, DataCollector.NameWithTypeParameters>,
         returnType: String,
         override: Boolean = false,
         final: Boolean = false,
@@ -74,7 +74,7 @@ abstract class AbstractVisitorGenerator(val referencesData: DataCollector.Refere
         printWithNoIndent(")")
     }
 
-    protected fun Printer.separatedOneLine(iterable: Iterable<Any>, separator: Any) {
+    private fun Printer.separatedOneLine(iterable: Iterable<Any>, separator: Any) {
         var first = true
         for (element in iterable) {
             if (!first) {
@@ -86,7 +86,7 @@ abstract class AbstractVisitorGenerator(val referencesData: DataCollector.Refere
         }
     }
 
-    protected fun Printer.generateDefaultImports() {
+    private fun Printer.generateDefaultImports() {
         referencesData.usedPackages.forEach {
             println("import ", it.asString(), ".*")
         }
@@ -118,8 +118,8 @@ abstract class AbstractVisitorGenerator(val referencesData: DataCollector.Refere
 
 
     fun allElementTypes() =
-        referencesData.back.let {
+        referencesData.back.let { map ->
-            it.keys + it.values.flatten()
+            map.keys + map.values.flatten()
         }.distinct()
 
     abstract fun Printer.generateContent()

compiler/fir/tree/visitors-generator/src/DataCollector.kt

@@ -11,9 +11,28 @@ import org.jetbrains.kotlin.psi.psiUtil.findDescendantOfType
 
 class DataCollector {
 
-    private val references = mutableMapOf<String, List<String>>()
+    data class NameWithTypeParameters(val name: String, val typeParameters: List<String>) : Comparable<NameWithTypeParameters> {
-    private val packagePerClass = mutableMapOf<String, FqName>()
+        override fun compareTo(other: NameWithTypeParameters): Int = name.compareTo(other.name)
-    private val baseTransformedTypes = mutableListOf<String>()
+
+        constructor(name: String, typeParameterString: String) :
+                this(name, typeParameterString.drop(1).dropLast(1).split(" ", "\t", ",").filter { it.isNotEmpty() })
+
+        constructor(name: String) : this(name, emptyList())
+
+        override fun toString(): String =
+            name + if (typeParameters.isEmpty()) "" else typeParameters.joinToString(prefix = "<", postfix = ">", separator = ", ")
+    }
+
+    private val references = mutableMapOf<NameWithTypeParameters, List<NameWithTypeParameters>>()
+    private val packagePerClass = mutableMapOf<NameWithTypeParameters, FqName>()
+    private val baseTransformedTypes = mutableListOf<NameWithTypeParameters>()
+
+    private fun KtSuperTypeListEntry.toNameWithTypeParameters(): NameWithTypeParameters? {
+        val type = typeAsUserType ?: return null
+        val name = type.referencedName ?: return null
+        val typeArguments = type.typeArgumentList?.text ?: ""
+        return NameWithTypeParameters(name, typeArguments)
+    }
 
     fun readFile(file: KtFile) {
         file.acceptChildren(object : KtVisitorVoid() {
@@ -24,14 +43,18 @@ class DataCollector {
             }
 
             override fun visitClass(klass: KtClass) {
-                val className = klass.name
+                val className = klass.name ?: run {
-                if (klass.isInterface() && className != null) {
+                    super.visitClass(klass)
-                    packagePerClass[className] = file.packageFqName
+                    return
+                }
+                val classNameWithParameters = NameWithTypeParameters(className, klass.typeParameterList?.text ?: "")
+                if (klass.isInterface()) {
+                    packagePerClass[classNameWithParameters] = file.packageFqName
                     val isBaseTT = klass.annotationEntries.any {
                         it.shortName?.asString() == BASE_TRANSFORMED_TYPE_ANNOTATION_NAME
                     }
                     if (isBaseTT) {
-                        baseTransformedTypes += className
+                        baseTransformedTypes += classNameWithParameters
                     }
 
                     val manual = klass.superTypeListEntries.find {
@@ -40,11 +63,11 @@ class DataCollector {
                         } != null
                     }
                     if (manual != null) {
-                        references[className] = listOfNotNull(manual.typeAsUserType?.referencedName)
+                        references[classNameWithParameters] = listOfNotNull(manual.toNameWithTypeParameters())
                     } else {
-                        references[className] =
+                        references[classNameWithParameters] =
                                 klass.superTypeListEntries.mapNotNull {
-                                    it.typeAsUserType?.referencedName
+                                    it.toNameWithTypeParameters()
                                 }
                     }
                 }
@@ -55,8 +78,8 @@ class DataCollector {
     }
 
 
-    private fun Map<String, List<String>>.computeBackReferences(): Map<String, List<String>> {
+    private fun <K> Map<K, List<K>>.computeBackReferences(): Map<K, List<K>> {
-        val result = mutableMapOf<String, List<String>>()
+        val result = mutableMapOf<K, List<K>>()
 
         this.forEach { (k, v) ->
             v.forEach {
@@ -66,7 +89,7 @@ class DataCollector {
         return result
     }
 
-    private fun Map<String, List<String>>.sorted(): Map<String, List<String>> {
+    private fun <K : Comparable<K>> Map<K, List<K>>.sortedMap(): Map<K, List<K>> {
         return this.toSortedMap().mapValues { (_, v) -> v.sorted() }
     }
 
@@ -80,17 +103,17 @@ class DataCollector {
 
         val cleanBack = back.filterKeys { it in keysToKeep }
         return ReferencesData(
-            cleanBack.computeBackReferences().sorted(),
+            cleanBack.computeBackReferences().sortedMap(),
-            cleanBack.sorted(),
+            cleanBack.sortedMap(),
             packagePerClass.filterKeys { it in keysToKeep }.values.distinct().sortedBy { it.asString() },
             baseTransformedTypes.sorted()
         )
     }
 
     data class ReferencesData(
-        val direct: Map<String, List<String>>,
+        val direct: Map<NameWithTypeParameters, List<NameWithTypeParameters>>,
-        val back: Map<String, List<String>>,
+        val back: Map<NameWithTypeParameters, List<NameWithTypeParameters>>,
         val usedPackages: List<FqName>,
-        val baseTransformedTypes: List<String>
+        val baseTransformedTypes: List<NameWithTypeParameters>
     )
 }

compiler/fir/tree/visitors-generator/src/ParametricTransformerGenerator.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.visitors.generator
 
+import org.jetbrains.kotlin.fir.visitors.generator.DataCollector.NameWithTypeParameters
 import org.jetbrains.kotlin.utils.Printer
 
 class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : AbstractVisitorGenerator(data) {
@@ -15,8 +16,8 @@ class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : Abstr
             generateFunction(
                 "transformElement",
                 mapOf(
-                    "element" to "E",
+                    "element" to NameWithTypeParameters("E"),
-                    "data" to "D"
+                    "data" to NameWithTypeParameters("D")
                 ),
                 constructCompositeBoxType("E"),
                 typeParameters = listOf("E : $FIR_ELEMENT_CLASS_NAME"),
@@ -40,7 +41,7 @@ class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : Abstr
         return "$TRANSFORMER_RESULT_NAME<$type>"
     }
 
-    val baseTypes = mutableMapOf<String, String>().also {
+    val baseTypes = mutableMapOf<NameWithTypeParameters, NameWithTypeParameters>().also {
         for (baseTransformedType in referencesData.baseTransformedTypes) {
             it[baseTransformedType] = baseTransformedType
         }
@@ -50,18 +51,19 @@ class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : Abstr
         }
     }
 
-    fun Printer.generateTrampolineVisit(className: String) {
+    fun Printer.generateTrampolineVisit(className: NameWithTypeParameters) {
-        val shortcutName = className.classNameWithoutFir
+        val shortcutName = className.name.classNameWithoutFir
         val parameterName = shortcutName.decapitalize().safeName
         generateFunction(
             name = "visit$shortcutName",
             parameters = mapOf(
                 parameterName to className,
-                "data" to "D"
+                "data" to NameWithTypeParameters("D")
             ),
-            returnType = constructCompositeBoxType(FIR_ELEMENT_CLASS_NAME),
+            returnType = constructCompositeBoxType(FIR_ELEMENT_CLASS_NAME.name),
             override = true,
-            final = true
+            final = true,
+            typeParameters = className.typeParameters
         ) {
             print("return ")
             generateCall("transform$shortcutName", listOf(parameterName, "data"))
@@ -69,23 +71,26 @@ class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : Abstr
         }
     }
 
-    fun Printer.generateTransformMethod(className: String, parent: String) {
+    fun Printer.generateTransformMethod(
+        className: NameWithTypeParameters,
+        parent: NameWithTypeParameters
+    ) {
         val baseType = baseTypes[className]
 
-        val shortcutName = className.classNameWithoutFir
+        val shortcutName = className.name.classNameWithoutFir
         val parameterName = shortcutName.decapitalize().safeName
         if (baseType == null) {
             generateFunction(
                 "transform$shortcutName",
                 mapOf(
-                    parameterName to "E",
+                    parameterName to NameWithTypeParameters("E"),
-                    "data" to "D"
+                    "data" to NameWithTypeParameters("D")
                 ),
                 constructCompositeBoxType("E"),
-                typeParameters = listOf("E : $parent")
+                typeParameters = listOf("E : $parent") + className.typeParameters
             ) {
                 print("return ")
-                generateCall("transform${parent.classNameWithoutFir}", listOf(parameterName, "data"))
+                generateCall("transform${parent.name.classNameWithoutFir}", listOf(parameterName, "data"))
                 printlnWithNoIndent()
             }
         } else {
@@ -93,12 +98,13 @@ class ParametricTransformerGenerator(data: DataCollector.ReferencesData) : Abstr
                 "transform$shortcutName",
                 mapOf(
                     parameterName to className,
-                    "data" to "D"
+                    "data" to NameWithTypeParameters("D")
                 ),
-                constructCompositeBoxType(baseType)
+                constructCompositeBoxType(baseType.name),
+                typeParameters = className.typeParameters
             ) {
                 print("return ")
-                generateCall("transform${parent.classNameWithoutFir}", listOf(parameterName, "data"))
+                generateCall("transform${parent.name.classNameWithoutFir}", listOf(parameterName, "data"))
                 printlnWithNoIndent()
             }
         }

compiler/fir/tree/visitors-generator/src/SimpleVisitorGenerator.kt

@@ -15,7 +15,7 @@ class SimpleVisitorGenerator(referencesData: DataCollector.ReferencesData) : Abs
                 "visitElement",
                 parameters = mapOf(
                     "element" to FIR_ELEMENT_CLASS_NAME,
-                    "data" to "D"
+                    "data" to DataCollector.NameWithTypeParameters("D")
                 ),
                 returnType = "R",
                 body = null
@@ -27,19 +27,23 @@ class SimpleVisitorGenerator(referencesData: DataCollector.ReferencesData) : Abs
         println("}")
     }
 
-    private fun Printer.generateVisit(className: String, parent: String) {
+    private fun Printer.generateVisit(
-        val shortcutName = className.classNameWithoutFir
+        className: DataCollector.NameWithTypeParameters,
+        parent: DataCollector.NameWithTypeParameters
+    ) {
+        val shortcutName = className.name.classNameWithoutFir
         val parameterName = shortcutName.decapitalize().safeName
         generateFunction(
             name = "visit$shortcutName",
             parameters = mapOf(
                 parameterName to className,
-                "data" to "D"
+                "data" to DataCollector.NameWithTypeParameters("D")
             ),
-            returnType = "R"
+            returnType = "R",
+            typeParameters = className.typeParameters
         ) {
             print("return ")
-            generateCall("visit${parent.classNameWithoutFir}", listOf(parameterName, "data"))
+            generateCall("visit${parent.name.classNameWithoutFir}", listOf(parameterName, "data"))
             println()
         }
     }

compiler/fir/tree/visitors-generator/src/UnitVisitorGenerator.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.visitors.generator
 
+import org.jetbrains.kotlin.fir.visitors.generator.DataCollector.NameWithTypeParameters
 import org.jetbrains.kotlin.utils.Printer
 
 class UnitVisitorGenerator(referencesData: DataCollector.ReferencesData) : AbstractVisitorGenerator(referencesData) {
@@ -33,34 +34,39 @@ class UnitVisitorGenerator(referencesData: DataCollector.ReferencesData) : Abstr
     }
 
 
-    private fun Printer.generateVisit(className: String, parent: String) {
+    private fun Printer.generateVisit(
-        val shortcutName = className.classNameWithoutFir
+        className: NameWithTypeParameters,
+        parent: NameWithTypeParameters
+    ) {
+        val shortcutName = className.name.classNameWithoutFir
         val parameterName = shortcutName.decapitalize().safeName
         generateFunction(
             name = "visit$shortcutName",
             parameters = mapOf(
                 parameterName to className
             ),
-            returnType = "Unit"
+            returnType = "Unit",
+            typeParameters = className.typeParameters
         ) {
             printIndent()
-            generateCall("visit${parent.classNameWithoutFir}", listOf(parameterName, "null"))
+            generateCall("visit${parent.name.classNameWithoutFir}", listOf(parameterName, "null"))
             println()
         }
     }
 
-    private fun Printer.generateTrampolineVisit(className: String) {
+    private fun Printer.generateTrampolineVisit(className: NameWithTypeParameters) {
-        val shortcutName = className.classNameWithoutFir
+        val shortcutName = className.name.classNameWithoutFir
         val parameterName = shortcutName.decapitalize().safeName
         generateFunction(
             name = "visit$shortcutName",
             parameters = mapOf(
                 parameterName to className,
-                "data" to "Nothing?"
+                "data" to NameWithTypeParameters("Nothing?")
             ),
             returnType = "Unit",
             override = true,
-            final = true
+            final = true,
+            typeParameters = className.typeParameters
         ) {
             printIndent()
             generateCall("visit$shortcutName", listOf(parameterName))

compiler/fir/tree/visitors-generator/src/VisitorsGenerator.kt

@@ -14,7 +14,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 
 
-const val FIR_ELEMENT_CLASS_NAME = "FirElement"
+val FIR_ELEMENT_CLASS_NAME = DataCollector.NameWithTypeParameters("FirElement")
 const val VISITOR_PACKAGE = "org.jetbrains.kotlin.fir.visitors"
 const val SIMPLE_VISITOR_NAME = "FirVisitor"
 const val UNIT_VISITOR_NAME = "FirVisitorVoid"
@@ -71,7 +71,10 @@ fun AbstractVisitorGenerator.runGenerator(file: File) {
 val String.classNameWithoutFir get() = this.removePrefix("Fir")
 
 
-fun DataCollector.ReferencesData.walkHierarchyTopDown(from: String, l: (p: String, e: String) -> Unit) {
+fun DataCollector.ReferencesData.walkHierarchyTopDown(
+    from: DataCollector.NameWithTypeParameters,
+    l: (p: DataCollector.NameWithTypeParameters, e: DataCollector.NameWithTypeParameters) -> Unit
+) {
     val referents = back[from] ?: return
     for (referent in referents) {
         l(from, referent)

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirTransformerGenerated.kt

@@ -138,6 +138,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformCall(delegatedConstructorCall, data)
     }
 
+    open fun <T> transformConstExpression(constExpression: FirConstExpression<T>, data: D): CompositeTransformResult<FirStatement> {
+        return transformExpression(constExpression, data)
+    }
+
     open fun transformErrorExpression(errorExpression: FirErrorExpression, data: D): CompositeTransformResult<FirStatement> {
         return transformExpression(errorExpression, data)
     }
@@ -310,6 +314,10 @@ abstract class FirTransformer<in D> : FirVisitor<CompositeTransformResult<FirEle
         return transformBody(body, data)
     }
 
+    final override fun <T> visitConstExpression(constExpression: FirConstExpression<T>, data: D): CompositeTransformResult<FirElement> {
+        return transformConstExpression(constExpression, data)
+    }
+
     final override fun visitErrorExpression(errorExpression: FirErrorExpression, data: D): CompositeTransformResult<FirElement> {
         return transformErrorExpression(errorExpression, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorGenerated.kt

@@ -138,6 +138,10 @@ abstract class FirVisitor<out R, in D> {
         return visitCall(delegatedConstructorCall, data)
     }
 
+    open fun <T> visitConstExpression(constExpression: FirConstExpression<T>, data: D): R {
+        return visitExpression(constExpression, data)
+    }
+
     open fun visitErrorExpression(errorExpression: FirErrorExpression, data: D): R {
         return visitExpression(errorExpression, data)
     }

compiler/fir/tree/visitors/org/jetbrains/kotlin/fir/visitors/FirVisitorVoidGenerated.kt

@@ -138,6 +138,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitCall(delegatedConstructorCall, null)
     }
 
+    open fun <T> visitConstExpression(constExpression: FirConstExpression<T>) {
+        visitExpression(constExpression, null)
+    }
+
     open fun visitErrorExpression(errorExpression: FirErrorExpression) {
         visitExpression(errorExpression, null)
     }
@@ -310,6 +314,10 @@ abstract class FirVisitorVoid : FirVisitor<Unit, Nothing?>() {
         visitBody(body)
     }
 
+    final override fun <T> visitConstExpression(constExpression: FirConstExpression<T>, data: Nothing?) {
+        visitConstExpression(constExpression)
+    }
+
     final override fun visitErrorExpression(errorExpression: FirErrorExpression, data: Nothing?) {
         visitErrorExpression(errorExpression)
     }