From 6dc75c2e5104f1bb6c64f34a68db677f1c6f8127 Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Tue, 8 Jun 2021 14:18:55 -0700
Subject: [PATCH] FIR IDE: allow KtClassOrObject to resolve to null KtSymbol

At PSI level, an enum entry is a `KtClassOrObject`. But the resolved
FIR counterpart `FirEnumEntry` is not an `FirRegularClass`.
---
 .../kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt | 5 +++--
 .../kotlin/idea/asJava/classes/firLightClassUtils.kt     | 4 ++--
 .../idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt | 9 ++++-----
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
index 6c1bd812df4..9f547dbd4d2 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtSymbolProvider.kt
@@ -42,7 +42,7 @@ public abstract class KtSymbolProvider : KtAnalysisSessionComponent() {
     public abstract fun getVariableSymbol(psi: KtProperty): KtVariableSymbol
     public abstract fun getAnonymousObjectSymbol(psi: KtObjectLiteralExpression): KtAnonymousObjectSymbol
     public abstract fun getClassOrObjectSymbol(psi: KtClassOrObject): KtClassOrObjectSymbol
-    public abstract fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol
+    public abstract fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol?
     public abstract fun getPropertyAccessorSymbol(psi: KtPropertyAccessor): KtPropertyAccessorSymbol
 
     public abstract fun getClassOrObjectSymbolByClassId(classId: ClassId): KtClassOrObjectSymbol?
@@ -102,7 +102,8 @@ public interface KtSymbolProviderMixIn : KtAnalysisSessionMixIn {
     public fun KtClassOrObject.getClassOrObjectSymbol(): KtClassOrObjectSymbol =
         analysisSession.symbolProvider.getClassOrObjectSymbol(this)
 
-    public fun KtClassOrObject.getNamedClassOrObjectSymbol(): KtNamedClassOrObjectSymbol =
+    /** Gets the corresponding class or object symbol or null if the given [KtClassOrObject] is an enum entry. */
+    public fun KtClassOrObject.getNamedClassOrObjectSymbol(): KtNamedClassOrObjectSymbol? =
         analysisSession.symbolProvider.getNamedClassOrObjectSymbol(this)
 
     public fun KtPropertyAccessor.getPropertyAccessorSymbol(): KtPropertyAccessorSymbol =
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
index 09e3fe27dd0..bee48e12919 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/asJava/classes/firLightClassUtils.kt
@@ -332,8 +332,8 @@ internal fun KtSymbolWithMembers.createInnerClasses(manager: PsiManager): List<F
 
 internal fun KtClassOrObject.checkIsInheritor(baseClassOrigin: KtClassOrObject, checkDeep: Boolean): Boolean {
     return analyseForLightClasses(this) {
-        val thisSymbol = this@checkIsInheritor.getNamedClassOrObjectSymbol()
-        val baseSymbol = baseClassOrigin.getNamedClassOrObjectSymbol()
+        val thisSymbol = this@checkIsInheritor.getNamedClassOrObjectSymbol() ?: return false
+        val baseSymbol = baseClassOrigin.getNamedClassOrObjectSymbol() ?: return false
 
         if (thisSymbol == baseSymbol) return@analyseForLightClasses false
 
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
index a5495afb0d9..5c618f3bc76 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirSymbolProvider.kt
@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirFile
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirOfType
-import org.jetbrains.kotlin.idea.fir.low.level.api.api.withFirDeclarationOfType
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
@@ -112,8 +109,10 @@ internal class KtFirSymbolProvider(
         }
     }
 
-    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol = withValidityAssertion {
+    override fun getNamedClassOrObjectSymbol(psi: KtClassOrObject): KtNamedClassOrObjectSymbol? = withValidityAssertion {
         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
         psi.withFirDeclarationOfType<FirRegularClass, KtNamedClassOrObjectSymbol>(resolveState) {
             firSymbolBuilder.classifierBuilder.buildNamedClassOrObjectSymbol(it)
         }