Optimization: make every FIR node extend FirPureAbstractElement

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

@@ -6,10 +6,10 @@
 package org.jetbrains.kotlin.fir
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirSession
+
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+abstract class FirPureAbstractElement : FirElement
 
 abstract class FirAbstractElement(
     final override val session: FirSession,
     final override val psi: PsiElement?
-) : FirElement
+) : FirPureAbstractElement()

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

@@ -11,7 +11,7 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirTransformer<D>, data: D) {
     val iterator = this.listIterator()
     while (iterator.hasNext()) {
-        val next = iterator.next()
+        val next = iterator.next() as FirPureAbstractElement
         val result = next.transform<T, D>(transformer, data)
         if (result.isSingle) {
             iterator.set(result.single)
@@ -26,6 +26,7 @@ fun <T : FirElement, D> MutableList<T>.transformInplace(transformer: FirTransfor
                 iterator.add(resultIterator.next())
             }
         }
+
     }
 }
 
@@ -55,7 +56,7 @@ fun <T : FirElement, D> MutableList<T>.transformInplaceWithBeforeOperation(
 }
 
 fun <T : FirElement, D> T.transformSingle(transformer: FirTransformer<D>, data: D): T {
-    return this.transform<T, D>(transformer, data).single
+    return (this as FirPureAbstractElement).transform<T, D>(transformer, data).single
 }
 
 fun ModuleInfo.dependenciesWithoutSelf(): Sequence<ModuleInfo> = dependencies().asSequence().filter { it != this }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirAnonymousInitializerImpl.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.declarations.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
@@ -14,10 +15,10 @@ import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 
 class FirAnonymousInitializerImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?,
+    psi: PsiElement?,
     override var body: FirBlock?
-) : FirAnonymousInitializer {
+) : FirAnonymousInitializer, FirAbstractElement(session, psi) {
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         body = body?.transformSingle(transformer, data)
         return this

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/impl/FirImportImpl.kt

@@ -6,15 +6,16 @@
 package org.jetbrains.kotlin.fir.declarations.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirImport
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
 class FirImportImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?,
+    psi: PsiElement?,
     override val importedFqName: FqName?,
     override val isAllUnder: Boolean,
     override val aliasName: Name?
-) : FirImport
+) : FirImport, FirAbstractElement(session, psi)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirAbstractAnnotatedTypeRef.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
@@ -14,10 +15,10 @@ import org.jetbrains.kotlin.fir.types.FirTypeRefWithNullability
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 
 abstract class FirAbstractAnnotatedTypeRef(
-    final override val session: FirSession,
+    session: FirSession,
-    final override val psi: PsiElement?,
+    psi: PsiElement?,
     final override val isMarkedNullable: Boolean
-) : FirTypeRefWithNullability {
+) : FirTypeRefWithNullability, FirAbstractElement(session, psi) {
     override val annotations = mutableListOf<FirAnnotationCall>()
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitBuiltinTypeRef.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTagImpl
@@ -15,10 +16,10 @@ import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.name.ClassId
 
 sealed class FirImplicitBuiltinTypeRef(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?,
+    psi: PsiElement?,
     val id: ClassId
-) : FirResolvedTypeRef {
+) : FirResolvedTypeRef, FirAbstractElement(session, psi) {
     override val annotations: List<FirAnnotationCall>
         get() = emptyList()
 

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirImplicitTypeRefImpl.kt

@@ -6,19 +6,21 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
+import org.jetbrains.kotlin.fir.FirPureAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 
 class FirImplicitTypeRefImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?
+    psi: PsiElement?
-) : FirImplicitTypeRef {
+) : FirImplicitTypeRef, FirAbstractElement(session, psi) {
     override val annotations: List<FirAnnotationCall>
         get() = emptyList()
 }
 
-object FirComputingImplicitTypeRef : FirImplicitTypeRef {
+object FirComputingImplicitTypeRef : FirImplicitTypeRef, FirPureAbstractElement() {
     override val psi: PsiElement?
         get() = null
     override val session: FirSession

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirResolvedFunctionTypeRefImpl.kt

@@ -6,27 +6,24 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.transformInplace
-import org.jetbrains.kotlin.fir.transformSingle
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirResolvedFunctionTypeRef
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 
 class FirResolvedFunctionTypeRefImpl(
-    override val psi: PsiElement?,
+    psi: PsiElement?,
-    override val session: FirSession,
+    session: FirSession,
     override val isMarkedNullable: Boolean,
     override val annotations: MutableList<FirAnnotationCall>,
     override var receiverTypeRef: FirTypeRef?,
     override val valueParameters: MutableList<FirValueParameter>,
     override var returnTypeRef: FirTypeRef,
     override val type: ConeKotlinType
-) : FirResolvedFunctionTypeRef {
+) : FirResolvedFunctionTypeRef, FirAbstractElement(session, psi) {
 
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         receiverTypeRef = receiverTypeRef?.transformSingle(transformer, data)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirResolvedTypeRefImpl.kt

@@ -6,14 +6,15 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 
 class FirResolvedTypeRefImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?,
+    psi: PsiElement?,
     override val type: ConeKotlinType,
     override val annotations: List<FirAnnotationCall> = emptyList()
-) : FirResolvedTypeRef
+) : FirResolvedTypeRef, FirAbstractElement(session, psi)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirStarProjectionImpl.kt

@@ -6,10 +6,11 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.types.FirStarProjection
 
 class FirStarProjectionImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?
+    psi: PsiElement?
-) : FirStarProjection
+) : FirStarProjection, FirAbstractElement(session, psi)

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/impl/FirTypeProjectionWithVarianceImpl.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.types.impl
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.fir.FirAbstractElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.transformSingle
@@ -15,11 +16,11 @@ import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.types.Variance
 
 class FirTypeProjectionWithVarianceImpl(
-    override val session: FirSession,
+    session: FirSession,
-    override val psi: PsiElement?,
+    psi: PsiElement?,
     override val variance: Variance,
     override var typeRef: FirTypeRef
-) : FirTypeProjectionWithVariance {
+) : FirTypeProjectionWithVariance, FirAbstractElement(session, psi) {
     override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirElement {
         typeRef = typeRef.transformSingle(transformer, data)