From 50247a034680cf71a85beeb2fe0113e00b2355e1 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Tue, 26 Jul 2022 22:00:23 +0200
Subject: [PATCH] [Analysis API FIR] fix getNamedClassOrObjectSymbol for MPP
 classes

---
 .../api/fir/symbols/KtFirSymbolProvider.kt    | 28 +++++++++----------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt
index 229ddde1b2d..67190ae5a9b 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt
@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbolOfT
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.firErrorWithAttachment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirSymbolAttachment
 import org.jetbrains.kotlin.analysis.utils.errors.withPsiAttachment
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.fullyExpandedClass
 import org.jetbrains.kotlin.fir.render
@@ -122,17 +123,7 @@ internal class KtFirSymbolProvider(
     }
 
     override fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol {
-        val firClass =
-            when (val firClassLike = psi.resolveToFirSymbolOfType<FirClassLikeSymbol<*>>(firResolveSession)) {
-                is FirTypeAliasSymbol -> firClassLike.fullyExpandedClass(firResolveSession.useSiteFirSession)
-                    ?: errorWithAttachment("${firClassLike.fir::class} should be expanded to the expected type alias") {
-                        withFirSymbolAttachment("firClassLikeSymbol", firClassLike)
-                        withPsiAttachment("ktClassOrObject", psi)
-                    }
-                is FirAnonymousObjectSymbol -> firClassLike
-                is FirRegularClassSymbol -> firClassLike
-            }
-
+        val firClass = psi.resolveToFirClassLikeSymbol()
         return firSymbolBuilder.classifierBuilder.buildClassOrObjectSymbol(firClass)
     }
 
@@ -140,10 +131,19 @@ internal class KtFirSymbolProvider(
         require(psi !is KtObjectDeclaration || psi.parent !is KtObjectLiteralExpression)
         // A KtClassOrObject may also map to an FirEnumEntry. Hence, we need to return null in this case.
         if (psi is KtEnumEntry) return null
-        return firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(
-            psi.resolveToFirSymbolOfType<FirRegularClassSymbol>(firResolveSession)
-        )
+        return firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(psi.resolveToFirClassLikeSymbol() as FirRegularClassSymbol)
+    }
 
+    private fun KtClassOrObject.resolveToFirClassLikeSymbol(): FirClassSymbol<*> {
+        return when (val firClassLike = resolveToFirSymbolOfType<FirClassLikeSymbol<*>>(firResolveSession)) {
+            is FirTypeAliasSymbol -> firClassLike.fullyExpandedClass(firResolveSession.useSiteFirSession)
+                ?: errorWithAttachment("${firClassLike.fir::class} should be expanded to the expected type alias") {
+                    withFirSymbolAttachment("firClassLikeSymbol", firClassLike)
+                    withPsiAttachment("ktClassOrObject", this@resolveToFirClassLikeSymbol)
+                }
+            is FirAnonymousObjectSymbol -> firClassLike
+            is FirRegularClassSymbol -> firClassLike
+        }
     }
 
     override fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol {