[FIR generator] Use TypeRef for denoting types

(instead of `Type` and `Importable`) This is a step towards commonizing the code generator between FIR and IR: KT-61970
2023-09-12 18:06:00 +02:00
parent 918cf183c1
commit 8798fdeb82
16 changed files with 113 additions and 110 deletions
@@ -16,9 +16,9 @@ import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeParameter
 import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeParameterRef
 import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeProjection
 import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.typeRef
-import org.jetbrains.kotlin.fir.tree.generator.FirTreeBuilder.valueParameter
 import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.fir.tree.generator.model.*
+import org.jetbrains.kotlin.generators.tree.TypeRef

 object FieldSets {
    val calleeReference by lazy { field("calleeReference", reference, withReplace = true) }
@@ -35,7 +35,7 @@ object FieldSets {

    val arguments by lazy { fieldList("arguments", expression) }

-    val declarations by lazy { fieldList(declaration.withArgs("E" to "*")) }
+    val declarations by lazy { fieldList(declaration.withArgs("E" to TypeRef.Star)) }

    val annotations by lazy {
        fieldList(
@@ -46,7 +46,7 @@ object FieldSets {
        ).withTransform(needTransformInOtherChildren = true)
    }

-    fun symbolWithPackage(packageName: String?, symbolClassName: String, argument: String? = null): Field {
+    fun symbolWithPackage(packageName: String, symbolClassName: String, argument: String? = null): Field {
        return field("symbol", type(packageName, symbolClassName), argument)
    }

@@ -34,6 +34,8 @@ import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder.Co
 import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.fir.tree.generator.model.*
 import org.jetbrains.kotlin.generators.tree.SimpleTypeArgument
+import org.jetbrains.kotlin.generators.tree.StandardTypes
+import org.jetbrains.kotlin.generators.tree.TypeRef
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource

 object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuilder) {
@@ -107,7 +109,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
            +field("deprecationsProvider", deprecationsProviderType).withReplace().apply { isMutable = true }
            +symbol("FirCallableSymbol", "out FirCallableDeclaration")

-            +field("containerSource", type(DeserializedContainerSource::class), nullable = true)
+            +field("containerSource", type<DeserializedContainerSource>(), nullable = true)
            +field("dispatchReceiverType", coneSimpleKotlinTypeType, nullable = true)

            +fieldList(contextReceiver, useMutableOrEmpty = true, withReplace = true)
@@ -169,7 +171,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
        }

        returnExpression.configure {
-            parentArg(jump, "E", function.withArgs("E" to "*"))
+            parentArg(jump, "E", function.withArgs("E" to TypeRef.Star))
            +field("result", expression).withTransform()
            needTransformOtherChildren()
        }
@@ -337,7 +339,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
        typeParameter.configure {
            +name
            +symbol("FirTypeParameterSymbol")
-            +field("containingDeclarationSymbol", firBasedSymbolType, "*").apply {
+            +field("containingDeclarationSymbol", firBasedSymbolType to listOf(TypeRef.Star)).apply {
                withBindThis = false
            }
            +field(varianceType)
@@ -547,7 +549,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
        }

        annotationArgumentMapping.configure {
-            +field("mapping", type("Map") to listOf(nameType, expression))
+            +field("mapping", StandardTypes.map to listOf(nameType, expression))
        }

        augmentedArraySetCall.configure {
@@ -746,7 +748,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild

        whenExpression.configure {
            +field("subject", expression, nullable = true).withTransform()
-            +field("subjectVariable", variable.withArgs("E" to "*"), nullable = true)
+            +field("subjectVariable", variable.withArgs("E" to TypeRef.Star), nullable = true)
            +fieldList("branches", whenBranch).withTransform()
            +field("exhaustivenessStatus", exhaustivenessStatusType, nullable = true, withReplace = true)
            +booleanField("usedAsExpression")
@@ -781,10 +783,11 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
    }
 }

-fun Element.withArgs(vararg replacements: Pair<String, String>): AbstractElement {
+// 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 ->
-        replaceMap[typeArgument.name]?.let { SimpleTypeArgument(it, null) } ?: typeArgument
+        replaceMap[typeArgument.name]?.let { SimpleTypeArgument(it.type, null) } ?: typeArgument
    }
    return ElementWithArguments(this, newArguments)
 }
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.ConeSimpleKotlinType
+import org.jetbrains.kotlin.generators.tree.TypeKind
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -26,30 +27,30 @@ import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.types.SmartcastStability
 import org.jetbrains.kotlin.types.Variance

-val sourceElementType = type(KtSourceElement::class)
-val sourceFileType = type(KtSourceFile::class)
-val sourceFileLinesMappingType = type(KtSourceFileLinesMapping::class)
+val sourceElementType = type<KtSourceElement>()
+val sourceFileType = type<KtSourceFile>()
+val sourceFileLinesMappingType = type<KtSourceFileLinesMapping>()
 val jumpTargetType = type("fir", "FirTarget")
 val constKindType = type("types", "ConstantValueKind")
 val operationType = type("fir.expressions", "FirOperation")
-val classKindType = type(ClassKind::class)
-val eventOccurrencesRangeType = type(EventOccurrencesRange::class)
+val classKindType = type<ClassKind>()
+val eventOccurrencesRangeType = type<EventOccurrencesRange>()
 val inlineStatusType = type("fir.declarations", "InlineStatus")
-val varianceType = type(Variance::class)
-val nameType = type(Name::class)
-val visibilityType = type(Visibility::class)
+val varianceType = type<Variance>()
+val nameType = type<Name>()
+val visibilityType = type<Visibility>()
 val effectiveVisibilityType = type("descriptors", "EffectiveVisibility")
-val modalityType = type(Modality::class)
-val smartcastStabilityType = type(SmartcastStability::class)
-val fqNameType = type(FqName::class)
-val classIdType = type(ClassId::class)
-val annotationUseSiteTargetType = type(AnnotationUseSiteTarget::class)
+val modalityType = type<Modality>()
+val smartcastStabilityType = type<SmartcastStability>()
+val fqNameType = type<FqName>()
+val classIdType = type<ClassId>()
+val annotationUseSiteTargetType = type<AnnotationUseSiteTarget>()
 val operationKindType = type("contracts.description", "LogicOperationKind")
-val coneKotlinTypeType = type(ConeKotlinType::class)
-val coneErrorTypeType = type(ConeErrorType::class)
-val coneSimpleKotlinTypeType = type(ConeSimpleKotlinType::class)
-val coneClassLikeTypeType = type(ConeClassLikeType::class)
-val standardClassIdsType = type(StandardClassIds::class)
+val coneKotlinTypeType = type<ConeKotlinType>()
+val coneErrorTypeType = type<ConeErrorType>()
+val coneSimpleKotlinTypeType = type<ConeSimpleKotlinType>()
+val coneClassLikeTypeType = type<ConeClassLikeType>()
+val standardClassIdsType = type<StandardClassIds>()

 val whenRefType = generatedType("", "FirExpressionRef<FirWhenExpression>")
 val referenceToSimpleExpressionType = generatedType("", "FirExpressionRef<FirExpression>")
@@ -84,7 +85,7 @@ val pureAbstractElementType = generatedType("FirPureAbstractElement")
 val coneContractElementType = type("fir.contracts.description", "ConeContractDescriptionElement")
 val coneEffectDeclarationType = type("fir.contracts.description", "ConeEffectDeclaration")
 val emptyContractDescriptionType = generatedType("contracts.impl", "FirEmptyContractDescription")
-val coneDiagnosticType = generatedType("diagnostics", "ConeDiagnostic")
+val coneDiagnosticType = generatedType("diagnostics", "ConeDiagnostic", kind = TypeKind.Interface)
 val coneStubDiagnosticType = generatedType("diagnostics", "ConeStubDiagnostic")

 val firImplementationDetailType = generatedType("FirImplementationDetail")
@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.fir.tree.generator.context

 import org.jetbrains.kotlin.fir.tree.generator.model.*
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
-import org.jetbrains.kotlin.generators.tree.Importable
-import org.jetbrains.kotlin.generators.tree.SimpleTypeArgument
-import org.jetbrains.kotlin.generators.tree.TypeArgumentWithMultipleUpperBounds
+import org.jetbrains.kotlin.generators.tree.*

 abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val builder: T) {
    inner class ConfigureContext(val element: Element) {
@@ -30,27 +27,27 @@ abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val
            }
        }

-        fun withArg(name: String, vararg upperBounds: String) {
-            withArg(name, upperBounds.map { Type(null, it) })
+        fun withArg(name: String) {
+            withArg(name, emptyList())
        }

-        fun withArg(name: String, upperBound: Importable, vararg upperBounds: Importable) {
+        fun withArg(name: String, upperBound: TypeRef, vararg upperBounds: TypeRef) {
            val allUpperBounds = mutableListOf(upperBound).apply { this += upperBounds }
            withArg(name, allUpperBounds)
        }

-        private fun withArg(name: String, upperBounds: List<Importable>) {
+        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: Importable) {
-            parentArg(parent, Type(null, argument), type)
+        fun parentArg(parent: Element, argument: String, type: TypeRef) {
+            parentArg(parent, NamedTypeParameterRef(argument), type)
        }

-        fun parentArg(parent: Element, argument: Importable, type: Importable) {
+        fun parentArg(parent: Element, argument: NamedTypeParameterRef, type: TypeRef) {
            require(parent in element.parents) {
                "$parent is not parent of $element"
            }
@@ -61,8 +58,8 @@ abstract class AbstractFieldConfigurator<T : AbstractFirTreeBuilder>(private val
            argMap[argument] = type
        }

-        fun Type.withArgs(vararg args: String): Pair<Type, List<Importable>> {
-            return this to args.map { Type(null, it) }
+        fun TypeRef.withArgs(vararg args: String): Pair<TypeRef, List<TypeRef>> {
+            return this to args.map { NamedTypeParameterRef(it) }
        }

        fun needTransformOtherChildren() {
@@ -7,9 +7,12 @@ package org.jetbrains.kotlin.fir.tree.generator.context

 import org.jetbrains.kotlin.fir.tree.generator.model.Element
 import org.jetbrains.kotlin.fir.tree.generator.model.IntermediateBuilder
-import org.jetbrains.kotlin.fir.tree.generator.model.Type
 import org.jetbrains.kotlin.fir.tree.generator.printer.BASE_PACKAGE
 import org.jetbrains.kotlin.utils.DummyDelegate
+import org.jetbrains.kotlin.generators.tree.ClassRef
+import org.jetbrains.kotlin.generators.tree.PositionTypeParameterRef
+import org.jetbrains.kotlin.generators.tree.TypeRef
+import org.jetbrains.kotlin.generators.tree.TypeKind
 import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
@@ -93,22 +96,21 @@ abstract class AbstractFirTreeBuilder {
    }
 }

-fun generatedType(type: String): Type = generatedType("", type)
+fun generatedType(type: String, kind: TypeKind = TypeKind.Class): ClassRef<PositionTypeParameterRef> = generatedType("", type, kind)

-fun generatedType(packageName: String, type: String): Type {
+fun generatedType(packageName: String, type: String, kind: TypeKind = TypeKind.Class): ClassRef<PositionTypeParameterRef> {
    val realPackage = BASE_PACKAGE + if (packageName.isNotBlank()) ".$packageName" else ""
-    return type(realPackage, type, exactPackage = true)
+    return type(realPackage, type, exactPackage = true, kind = kind)
 }

-fun type(packageName: String?, type: String, exactPackage: Boolean = false): Type {
-    val realPackage = if (exactPackage) packageName else packageName?.let { "org.jetbrains.kotlin.$it" }
-    return Type(realPackage, type)
+fun type(
+    packageName: String,
+    type: String,
+    exactPackage: Boolean = false,
+    kind: TypeKind = TypeKind.Interface,
+): ClassRef<PositionTypeParameterRef> {
+    val realPackage = if (exactPackage) packageName else packageName.let { "org.jetbrains.kotlin.$it" }
+    return org.jetbrains.kotlin.generators.tree.type(realPackage, type, kind)
 }

-fun type(type: String): Type = type(null, type)
-
-fun type(klass: KClass<*>): Type {
-    val fqn = klass.qualifiedName!!
-    val name = klass.simpleName!!
-    return type(fqn.dropLast(name.length + 1), name, exactPackage = true)
-}
+inline fun <reified T : Any> type() = org.jetbrains.kotlin.generators.tree.type<T>()
@@ -7,10 +7,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.util.set
-import org.jetbrains.kotlin.generators.tree.ImplementationKind
-import org.jetbrains.kotlin.generators.tree.Importable
-import org.jetbrains.kotlin.generators.tree.TypeArgument
-import org.jetbrains.kotlin.generators.tree.typeWithArguments
+import org.jetbrains.kotlin.generators.tree.*
 import org.jetbrains.kotlin.generators.tree.AbstractElement as CommonAbstractElement

 interface AbstractElement : CommonAbstractElement<AbstractElement, Field> {
@@ -43,7 +40,7 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
    override var defaultImplementation: Implementation? = null
    override val customImplementations = mutableListOf<Implementation>()
    override val typeArguments = mutableListOf<TypeArgument>()
-    override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<Importable, Importable>>()
+    override val parentsArguments = mutableMapOf<AbstractElement, MutableMap<TypeRef, TypeRef>>()
    override var kind: ImplementationKind? = null
        set(value) {
            if (value !in allowedKinds) {
@@ -61,7 +58,7 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
    override var doesNotNeedImplementation: Boolean = false

    override val needTransformOtherChildren: Boolean get() = _needTransformOtherChildren || parents.any { it.needTransformOtherChildren }
-    override val overridenFields: MutableMap<Field, MutableMap<Importable, Boolean>> = mutableMapOf()
+    override val overridenFields: MutableMap<Field, MutableMap<Field, Boolean>> = mutableMapOf()
    override val useNullableForReplace: MutableSet<Field> = mutableSetOf()
    override val allImplementations: List<Implementation> by lazy {
        if (doesNotNeedImplementation) {
@@ -87,7 +84,7 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
                existingField.withReplace = parentField.withReplace || existingField.withReplace
                existingField.parentHasSeparateTransform = parentField.needsSeparateTransform
                if (parentField.type != existingField.type && parentField.withReplace) {
-                    existingField.overridenTypes += parentField
+                    existingField.overridenTypes += parentField.typeRef
                    overridenFields[existingField, parentField] = false
                } else {
                    overridenFields[existingField, parentField] = true
@@ -107,8 +104,8 @@ class Element(override val name: String, kind: Kind) : AbstractElement {
        parents.forEach { parent ->
            val fields = parent.allFields.map { field ->
                val copy = (field as? SimpleField)?.let { simpleField ->
-                    parentsArguments[parent]?.get(Type(null, simpleField.type))?.let {
-                        simpleField.replaceType(Type(it.packageName, it.type))
+                    parentsArguments[parent]?.get(NamedTypeParameterRef(simpleField.type))?.let {
+                        simpleField.replaceType(it)
                    }
                } ?: field.copy()
                copy.apply {
@@ -9,25 +9,27 @@ import org.jetbrains.kotlin.fir.tree.generator.context.AbstractFirTreeBuilder
 import org.jetbrains.kotlin.fir.tree.generator.context.type
 import org.jetbrains.kotlin.generators.tree.typeWithArguments
 import org.jetbrains.kotlin.generators.tree.Importable
+import org.jetbrains.kotlin.generators.tree.NamedTypeParameterRef
+import org.jetbrains.kotlin.generators.tree.TypeRef

 // ----------- Simple field -----------

-fun field(name: String, type: String, packageName: String?, customType: Importable? = null, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(name: String, type: String, packageName: String?, customType: TypeRef? = null, nullable: Boolean = false, withReplace: Boolean = false): Field {
    return SimpleField(name, type, packageName, customType, nullable, withReplace)
 }

-fun field(name: String, type: Type, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(name: String, type: TypeRef, nullable: Boolean = false, withReplace: Boolean = false): Field {
    return SimpleField(name, type.typeWithArguments, type.packageName, null, nullable, withReplace)
 }

-fun field(name: String, typeWithArgs: Pair<Type, List<Importable>>, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(name: String, typeWithArgs: Pair<TypeRef, List<TypeRef>>, nullable: Boolean = false, withReplace: Boolean = false): Field {
    val (type, args) = typeWithArgs
    return SimpleField(name, type.typeWithArguments, type.packageName, null, nullable, withReplace).apply {
        arguments += args
    }
 }

-fun field(type: Type, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(type: TypeRef, nullable: Boolean = false, withReplace: Boolean = false): Field {
    return SimpleField(type.type.replaceFirstChar(Char::lowercaseChar), type.typeWithArguments, type.packageName, null, nullable, withReplace)
 }

@@ -45,11 +47,11 @@ fun intField(name: String, withReplace: Boolean = false): Field {

 // ----------- Fir field -----------

-fun field(name: String, type: Type, argument: String? = null, nullable: Boolean = false, withReplace: Boolean = false): Field {
+fun field(name: String, type: TypeRef, argument: String? = null, nullable: Boolean = false, withReplace: Boolean = false): Field {
    return if (argument == null) {
        field(name, type, nullable, withReplace)
    } else {
-        field(name, type to listOf(type(argument)), nullable, withReplace)
+        field(name, type to listOf(NamedTypeParameterRef(argument)), nullable, withReplace)
    }
 }

@@ -63,7 +65,7 @@ fun field(element: Element, nullable: Boolean = false, withReplace: Boolean = fa

 // ----------- Field list -----------

-fun fieldList(name: String, type: Importable, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false): Field {
+fun fieldList(name: String, type: TypeRef, withReplace: Boolean = false, useMutableOrEmpty: Boolean = false): Field {
    return FieldList(name, type, withReplace, useMutableOrEmpty)
 }

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.fir.tree.generator.model

 import org.jetbrains.kotlin.fir.tree.generator.printer.generics
-import org.jetbrains.kotlin.generators.tree.typeWithArguments
-import org.jetbrains.kotlin.generators.tree.ArbitraryImportable
-import org.jetbrains.kotlin.generators.tree.Importable
-import org.jetbrains.kotlin.generators.tree.AbstractField
+import org.jetbrains.kotlin.generators.tree.*

 sealed class Field : AbstractField() {
    open var withReplace: Boolean = false
@@ -25,7 +22,7 @@ sealed class Field : AbstractField() {
    open var isMutableInInterface: Boolean = false
    open val fromDelegate: Boolean get() = false

-    open val overridenTypes: MutableSet<Importable> = mutableSetOf()
+    open val overridenTypes: MutableSet<TypeRef> = mutableSetOf()
    open var useNullableForReplace: Boolean = false
    open var notNull: Boolean = false

@@ -78,7 +75,7 @@ class FieldWithDefault(val origin: Field) : Field() {
        get() = origin.needTransformInOtherChildren
        set(_) {}

-    override val arguments: MutableList<Importable>
+    override val arguments: MutableList<TypeRef>
        get() = origin.arguments

    override val fullQualifiedName: String?
@@ -113,7 +110,7 @@ class FieldWithDefault(val origin: Field) : Field() {
    override var customSetter: String? = null
    override var fromDelegate: Boolean = false
    var needAcceptAndTransform: Boolean = true
-    override val overridenTypes: MutableSet<Importable>
+    override val overridenTypes: MutableSet<TypeRef>
        get() = origin.overridenTypes

    override val arbitraryImportables: MutableList<Importable>
@@ -138,7 +135,7 @@ class SimpleField(
    override val name: String,
    override val type: String,
    override val packageName: String?,
-    val customType: Importable? = null,
+    val customType: TypeRef? = null,
    override val nullable: Boolean,
    override var withReplace: Boolean,
    override var isVolatile: Boolean = false,
@@ -170,7 +167,7 @@ class SimpleField(
        }
    }

-    fun replaceType(newType: Type) = SimpleField(
+    fun replaceType(newType: TypeRef) = SimpleField(
        name = name,
        type = newType.type,
        packageName = newType.packageName,
@@ -195,7 +192,7 @@ class FirField(
 ) : Field() {
    init {
        if (element is ElementWithArguments) {
-            arguments += element.typeArguments.map { Type(null, it.name) }
+            arguments += element.typeArguments.map { NamedTypeParameterRef(it.name) }
        }
    }

@@ -229,7 +226,7 @@ class FirField(

 class FieldList(
    override val name: String,
-    val baseType: Importable,
+    val baseType: TypeRef,
    override var withReplace: Boolean,
    useMutableOrEmpty: Boolean = false
 ) : Field() {
@@ -1,12 +0,0 @@
-/*
- * Copyright 2010-2019 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.fir.tree.generator.model
-
-import org.jetbrains.kotlin.generators.tree.Importable
-
-data class Type(override val packageName: String?, override val type: String) : Importable {
-    override fun getTypeWithArguments(notNull: Boolean): String = type
-}
@@ -12,6 +12,7 @@ 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.typeWithArguments
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.withIndent
@@ -101,7 +102,7 @@ fun SmartPrinter.printElement(element: Element) {
                println("transformer.transform$name(this, data) as E")
            }

-            fun Field.replaceDeclaration(override: Boolean, overridenType: Importable? = null, forceNullable: Boolean = false) {
+            fun Field.replaceDeclaration(override: Boolean, overridenType: TypeRef? = null, forceNullable: Boolean = false) {
                println()
                if (name == "source") {
                    println("@FirImplementationDetail")
@@ -9,6 +9,7 @@ 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.TypeRef
 import org.jetbrains.kotlin.generators.tree.typeWithArguments
 import org.jetbrains.kotlin.utils.SmartPrinter
 import org.jetbrains.kotlin.utils.addToStdlib.ifNotEmpty
@@ -331,7 +332,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {

            fun generateReplace(
                field: Field,
-                overridenType: Importable? = null,
+                overridenType: TypeRef? = null,
                forceNullable: Boolean = false,
                body: () -> Unit,
            ) {
@@ -53,8 +53,8 @@ fun Implementation.collectImports(base: List<String> = emptyList(), kind: Import
                + arbitraryImportables.mapNotNull { it.fullQualifiedName }
                + parents.mapNotNull { it.fullQualifiedName }
                + listOfNotNull(
-            pureAbstractElementType.fullQualifiedName?.takeIf { needPureAbstractElement },
-            firImplementationDetailType.fullQualifiedName?.takeIf { isPublic || requiresOptIn },
+            pureAbstractElementType.fullQualifiedName.takeIf { needPureAbstractElement },
+            firImplementationDetailType.fullQualifiedName.takeIf { isPublic || requiresOptIn },
        ),
        kind,
    )
@@ -65,7 +65,7 @@ fun Element.collectImports(): List<String> {
    baseTypes += AbstractFirTreeBuilder.baseFirElement.fullQualifiedName
    baseTypes += parentsArguments.values.flatMap { it.values }.mapNotNull { it.fullQualifiedName }
    if (needPureAbstractElement) {
-        baseTypes += pureAbstractElementType.fullQualifiedName!!
+        baseTypes += pureAbstractElementType.fullQualifiedName
    }
    return collectImportsInternal(
        baseTypes,
@@ -107,7 +107,7 @@ private fun List<String>.filterRedundantImports(
 ): List<String> {
    val realPackageName = "$packageName.${kind.postfix}"
    return filter { fqn ->
-        fqn.dropLastWhile { it != '.' } != realPackageName
+        !fqn.startsWith("kotlin.") && fqn.dropLastWhile { it != '.' } != realPackageName
    }.distinct().sorted() + "$VISITOR_PACKAGE.*"
 }

@@ -128,7 +128,7 @@ fun transformFunctionDeclaration(transformName: String, returnType: String): Str
    return "fun <D> transform$transformName(transformer: FirTransformer<D>, data: D): $returnType"
 }

-fun Field.replaceFunctionDeclaration(overridenType: Importable? = null, forceNullable: Boolean = false): String {
+fun Field.replaceFunctionDeclaration(overridenType: TypeRef? = null, forceNullable: Boolean = false): String {
    val capName = name.replaceFirstChar(Char::uppercaseChar)
    val type = overridenType?.typeWithArguments ?: typeWithArguments

@@ -8,7 +8,8 @@ package org.jetbrains.kotlin.generators.tree
 /**
 * A common interface representing a FIR or IR tree element.
 */
-interface AbstractElement<Element : AbstractElement<Element, Field>, Field : AbstractField> : FieldContainer, ImplementationKindOwner {
+interface AbstractElement<Element : AbstractElement<Element, Field>, Field : AbstractField> : FieldContainer, ImplementationKindOwner,
+    TypeRef /* TODO: Replace with ElementOrRef */ {

    val name: String

@@ -18,9 +19,9 @@ interface AbstractElement<Element : AbstractElement<Element, Field>, Field : Abs

    val typeArguments: List<TypeArgument>

-    val parentsArguments: Map<Element, Map<Importable, Importable>>
+    val parentsArguments: Map<Element, Map<TypeRef, TypeRef>>

-    val overridenFields: Map<Field, Map<Importable, Boolean>>
+    val overridenFields: Map<Field, Map<Field, Boolean>>

    val isSealed: Boolean
        get() = false
@@ -9,7 +9,10 @@ abstract class AbstractField : Importable {

    abstract val name: String

-    open val arguments = mutableListOf<Importable>()
+    val typeRef: TypeRef
+        get() = type(packageName!!, type).withArgs(*arguments.toTypedArray()).copy(nullable = nullable)
+
+    open val arguments = mutableListOf<TypeRef>()

    abstract val nullable: Boolean

@@ -0,0 +1,10 @@
+/*
+ * 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
+
+object StandardTypes {
+    val map = type<Map<*, *>>()
+}
@@ -6,11 +6,11 @@
 package org.jetbrains.kotlin.generators.tree

 sealed class TypeArgument(val name: String) {
-    abstract val upperBounds: List<Importable>
+    abstract val upperBounds: List<TypeRef>
 }

-class SimpleTypeArgument(name: String, val upperBound: Importable?) : TypeArgument(name) {
-    override val upperBounds: List<Importable> = listOfNotNull(upperBound)
+class SimpleTypeArgument(name: String, val upperBound: TypeRef?) : TypeArgument(name) {
+    override val upperBounds: List<TypeRef> = listOfNotNull(upperBound)

    override fun toString(): String {
        var result = name
@@ -21,7 +21,7 @@ class SimpleTypeArgument(name: String, val upperBound: Importable?) : TypeArgume
    }
 }

-class TypeArgumentWithMultipleUpperBounds(name: String, override val upperBounds: List<Importable>) : TypeArgument(name) {
+class TypeArgumentWithMultipleUpperBounds(name: String, override val upperBounds: List<TypeRef>) : TypeArgument(name) {
    override fun toString(): String {
        return name
    }