Refactoring: get rid of firSafeNullable, enhance getSymbolByLookupTag

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt

@@ -23,17 +23,19 @@ abstract class FirSymbolProvider : FirSessionComponent {
 
     abstract fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>?
 
+    fun getSymbolByLookupTag(lookupTag: ConeClassLikeLookupTag): FirClassLikeSymbol<*>? {
+        (lookupTag as? ConeClassLikeLookupTagImpl)
+            ?.boundSymbol?.takeIf { it.first === this }?.let { return it.second }
+
+        return getClassLikeSymbolByFqName(lookupTag.classId).also {
+            (lookupTag as? ConeClassLikeLookupTagImpl)
+                ?.boundSymbol = Pair(this, it)
+        }
+    }
+
     fun getSymbolByLookupTag(lookupTag: ConeClassifierLookupTag): FirClassifierSymbol<*>? {
         return when (lookupTag) {
-            is ConeClassLikeLookupTag -> {
-                (lookupTag as? ConeClassLikeLookupTagImpl)
-                    ?.boundSymbol?.takeIf { it.first === this }?.let { return it.second }
-
-                getClassLikeSymbolByFqName(lookupTag.classId).also {
-                    (lookupTag as? ConeClassLikeLookupTagImpl)
-                        ?.boundSymbol = Pair(this, it)
-                }
-            }
+            is ConeClassLikeLookupTag -> getSymbolByLookupTag(lookupTag)
             is ConeClassifierLookupTagWithFixedSymbol -> lookupTag.symbol
             else -> error("Unknown lookupTag type: ${lookupTag::class}")
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -38,7 +38,7 @@ inline fun <K, V, VA : V> MutableMap<K, V>.getOrPut(key: K, defaultValue: (K) ->
 val FirSession.firSymbolProvider: FirSymbolProvider by componentArrayAccessor()
 val FirSession.correspondingSupertypesCache: FirCorrespondingSupertypesCache by componentArrayAccessor()
 
-fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassifierSymbol<*>? {
+fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): FirClassLikeSymbol<*>? {
     val firSymbolProvider = useSiteSession.firSymbolProvider
     return firSymbolProvider.getSymbolByLookupTag(this)
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ScopeUtils.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.resolve.transformers.firSafeNullable
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirCompositeScope
@@ -23,9 +22,8 @@ fun ConeKotlinType.scope(useSiteSession: FirSession, scopeSession: ScopeSession)
         is ConeClassErrorType -> null
         is ConeAbbreviatedType -> directExpansionType(useSiteSession)?.scope(useSiteSession, scopeSession)
         is ConeClassLikeType -> {
-            // For ConeClassLikeType they might be a type alias instead of a regular class
-            // TODO: support that case and switch back to `firUnsafe` instead of `firSafeNullable`
-            val fir = this.lookupTag.toSymbol(useSiteSession)?.firSafeNullable<FirRegularClass>() ?: return null
+            // TODO: for ConeClassLikeType they might be a type alias instead of a regular class
+            val fir = this.lookupTag.toSymbol(useSiteSession)?.fir as? FirRegularClass ?: return null
             wrapSubstitutionScopeIfNeed(useSiteSession, fir.buildUseSiteScope(useSiteSession, scopeSession)!!, scopeSession)
         }
         is ConeTypeParameterType -> {
@@ -38,8 +36,12 @@ fun ConeKotlinType.scope(useSiteSession: FirSession, scopeSession: ScopeSession)
             )
         }
         is ConeFlexibleType -> lowerBound.scope(useSiteSession, scopeSession)
-        is ConeIntersectionType -> FirCompositeScope(intersectedTypes.mapNotNullTo(mutableListOf()) { it.scope(useSiteSession, scopeSession) })
-        else -> error("Failed type ${this}")
+        is ConeIntersectionType -> FirCompositeScope(
+            intersectedTypes.mapNotNullTo(mutableListOf()) {
+                it.scope(useSiteSession, scopeSession)
+            }
+        )
+        else -> error("Failed type $this")
     }
 }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyExpressionBlock
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
 import org.jetbrains.kotlin.fir.resolve.withNullability
 import org.jetbrains.kotlin.fir.returnExpressions

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt

@@ -1037,10 +1037,6 @@ inline fun <reified T : FirElement> FirBasedSymbol<*>.firUnsafe(): T {
     return fir
 }
 
-inline fun <reified T : FirElement> FirBasedSymbol<*>.firSafeNullable(): T? {
-    return fir as? T
-}
-
 internal inline var FirExpression.resultType: FirTypeRef
     get() = typeRef
     set(type) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/impl/FirAbstractImportingScope.kt

@@ -33,7 +33,7 @@ abstract class FirAbstractImportingScope(
         val symbol = provider.getClassLikeSymbolByFqName(classId) ?: error("No scope/symbol for $classId")
         if (symbol is FirTypeAliasSymbol) {
             val expansionSymbol = symbol.fir.expandedConeType?.lookupTag?.toSymbol(session)
-            if (expansionSymbol is FirClassLikeSymbol<*>) {
+            if (expansionSymbol != null) {
                 return getStaticsScope(expansionSymbol.classId)
             }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/FirCorrespondingSupertypesCache.kt

@@ -26,7 +26,7 @@ class FirCorrespondingSupertypesCache(private val session: FirSession) : FirSess
     ): List<ConeClassLikeType>? {
         if (type !is ConeClassLikeType || supertypeConstructor !is FirClassLikeSymbol<*>) return null
 
-        val symbol = type.lookupTag.toSymbol(session) as? FirClassLikeSymbol<*> ?: return null
+        val symbol = type.lookupTag.toSymbol(session) ?: return null
         if (symbol == supertypeConstructor) return listOf(captureType(type))
 
         if (symbol !in cache) {

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 interface FirResolvedCallableReference : FirNamedReference {
+    // Can't write FirCallableSymbol here because of enums
     val resolvedSymbol: AbstractFirBasedSymbol<*>
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R =

compiler/fir/tree/src/org/jetbrains/kotlin/fir/symbols/impl/ConeClassLikeLookupTagImpl.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
 import org.jetbrains.kotlin.name.ClassId
 
 class ConeClassLikeLookupTagImpl(override val classId: ClassId) : ConeClassLikeLookupTag() {
-    var boundSymbol: Pair<*, FirClassifierSymbol<*>?>? = null
+    var boundSymbol: Pair<*, FirClassLikeSymbol<*>?>? = null
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -13,12 +13,10 @@ import org.jetbrains.kotlin.fir.FirNamedReference
 import org.jetbrains.kotlin.fir.FirResolvedCallableReference
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirAbstractMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.impl.FirClassImpl
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.references.FirErrorNamedReference
 import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.directExpansionType
-import org.jetbrains.kotlin.fir.resolve.transformers.firSafeNullable
 import org.jetbrains.kotlin.fir.resolve.transformers.firUnsafe
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
@@ -307,7 +305,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
                 if (it is FirResolvedQualifier) {
                     val lookupTag = (it.typeRef as FirResolvedTypeRefImpl).coneTypeSafe<ConeClassType>()?.lookupTag
                     val type = lookupTag?.let {
-                        symbolProvider.getSymbolByLookupTag(it)?.firSafeNullable<FirClass>()?.superTypeRefs?.first()?.render()
+                        (symbolProvider.getSymbolByLookupTag(it)?.fir as? FirClass)?.superTypeRefs?.first()?.render()
                     }
                     if (type != null) return@joinTo type
                 }
@@ -382,9 +380,9 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
                 data.append(coneClassType.lookupTag.classId.getWithoutCurrentPackage())
 
-                val typeParameters = symbolProvider.getSymbolByLookupTag(coneClassType.lookupTag)
-                    ?.firSafeNullable<FirClassImpl>()
-                    ?.typeParameters ?: listOf<FirTypeParameter>()
+                val typeParameters =
+                    (symbolProvider.getSymbolByLookupTag(coneClassType.lookupTag)?.fir as? FirRegularClass)
+                        ?.typeParameters ?: emptyList()
                 renderListInTriangles(typeParameters, data)
                 visitArguments(delegatedConstructorCall.arguments, data)
             } else {
@@ -408,7 +406,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
                         visitArguments(functionCall.arguments, data)
                     } else {
                         data.append("fun ")
-                        val firFunction = callee.resolvedSymbol.firSafeNullable<FirAbstractMemberDeclaration>()
+                        val firFunction = callee.resolvedSymbol.fir as? FirAbstractMemberDeclaration
                         firFunction?.let { renderListInTriangles(it.typeParameters, data, true) }
 
                         data.append(renderSymbol(callee.resolvedSymbol))