[FIR IDE] Supporting member scopes in EnumEntries

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.idea.frontend.api.components.*
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithKind
 import org.jetbrains.kotlin.idea.frontend.api.symbols.pointers.KtSymbolPointer
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
@@ -69,9 +70,9 @@ abstract class KtAnalysisSession(override val token: ValidityToken) : ValidityTo
 
     fun KtSymbolWithKind.getContainingSymbol(): KtSymbolWithKind? = containingDeclarationProvider.getContainingDeclaration(this)
 
-    fun KtClassOrObjectSymbol.getMemberScope(): KtMemberScope = scopeProvider.getMemberScope(this)
+    fun KtSymbolWithDeclarations.getMemberScope(): KtMemberScope = scopeProvider.getMemberScope(this)
 
-    fun KtClassOrObjectSymbol.getDeclaredMemberScope(): KtDeclaredMemberScope = scopeProvider.getDeclaredMemberScope(this)
+    fun KtSymbolWithDeclarations.getDeclaredMemberScope(): KtDeclaredMemberScope = scopeProvider.getDeclaredMemberScope(this)
 
     fun KtPackageSymbol.getPackageScope(): KtPackageScope = scopeProvider.getPackageScope(this)
 

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtScopeProvider.kt

@@ -5,17 +5,16 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.components
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KtScopeProvider : KtAnalysisSessionComponent() {
-    abstract fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope
+    abstract fun getMemberScope(classSymbol: KtSymbolWithDeclarations): KtMemberScope
-    abstract fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope
+    abstract fun getDeclaredMemberScope(classSymbol: KtSymbolWithDeclarations): KtDeclaredMemberScope
     abstract fun getPackageScope(packageSymbol: KtPackageSymbol): KtPackageScope
     abstract fun getCompositeScope(subScopes: List<KtScope>): KtCompositeScope
 

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/scopes/KtScope.kt

@@ -5,8 +5,10 @@
 
 package org.jetbrains.kotlin.idea.frontend.api.scopes
 
+import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -42,11 +44,11 @@ interface KtCompositeScope : KtScope {
 }
 
 interface KtMemberScope : KtScope {
-    val owner: KtClassOrObjectSymbol
+    val owner: KtSymbolWithDeclarations
 }
 
 interface KtDeclaredMemberScope : KtScope {
-    val owner: KtClassOrObjectSymbol
+    val owner: KtSymbolWithDeclarations
 }
 
 interface KtPackageScope : KtScope, KtSubstitutedScope<KtPackageScope> {

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtClassLikeSymbol.kt

@@ -35,7 +35,8 @@ abstract class KtClassOrObjectSymbol : KtClassLikeSymbol(),
     KtSymbolWithTypeParameters,
     KtSymbolWithModality<KtSymbolModality>,
     KtSymbolWithVisibility,
-    KtAnnotatedSymbol {
+    KtAnnotatedSymbol,
+    KtSymbolWithDeclarations {
     abstract val classKind: KtClassKind
 
     abstract val isInner: Boolean

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/KtVariableLikeSymbol.kt

@@ -14,10 +14,12 @@ sealed class KtVariableLikeSymbol : KtCallableSymbol(), KtTypedSymbol, KtNamedSy
     abstract override fun createPointer(): KtSymbolPointer<KtVariableLikeSymbol>
 }
 
-abstract class KtEnumEntrySymbol : KtVariableLikeSymbol(), KtSymbolWithKind {
+abstract class KtEnumEntrySymbol : KtVariableLikeSymbol(), KtSymbolWithDeclarations, KtSymbolWithKind {
     final override val symbolKind: KtSymbolKind get() = KtSymbolKind.MEMBER
     abstract val containingEnumClassIdIfNonLocal: ClassId?
 
+    abstract val hasBody: Boolean
+
     abstract override fun createPointer(): KtSymbolPointer<KtEnumEntrySymbol>
 }
 

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/markers/KtAnnotatedSymbol.kt

@@ -17,10 +17,6 @@ abstract class KtAnnotationCall {
     abstract val arguments: List<KtNamedConstantValue>
 }
 
-sealed class KtConstantValue
-object KtUnsupportedConstantValue : KtConstantValue()
-
-data class KtSimpleConstantValue<T>(val constant: T) : KtConstantValue()
 data class KtNamedConstantValue(val name: String, val expression: KtConstantValue)
 
 interface KtAnnotatedSymbol : KtSymbol {

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/markers/KtConstantValue.kt

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2010-2020 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.idea.frontend.api.symbols.markers
+
+sealed class KtConstantValue
+object KtUnsupportedConstantValue : KtConstantValue()
+
+data class KtSimpleConstantValue<T>(val constant: T?) : KtConstantValue()

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/markers/KtSymbolWithDeclarations.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2010-2020 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.idea.frontend.api.symbols.markers
+
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+
+interface KtSymbolWithDeclarations : KtSymbol

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/symbols/markers/KtSymbolWithModality.kt

@@ -17,4 +17,5 @@ sealed class KtCommonSymbolModality : KtSymbolModality() {
     object FINAL : KtCommonSymbolModality()
     object ABSTRACT : KtCommonSymbolModality()
     object OPEN : KtCommonSymbolModality()
+    object UNKNOWN : KtCommonSymbolModality()
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirScopeProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.components
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.fir.declarations.FirClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.scope
 import org.jetbrains.kotlin.fir.scopes.FakeOverrideTypeCalculator
@@ -25,13 +26,15 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.scopes.*
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirEnumEntrySymbol
 import org.jetbrains.kotlin.idea.frontend.api.fir.types.KtFirType
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.EnclosingDeclarationContext
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.buildCompletionContext
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.*
-import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassOrObjectSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtPackageSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
 import org.jetbrains.kotlin.psi.KtElement
@@ -50,15 +53,24 @@ internal class KtFirScopeProvider(
     private val firResolveState by weakRef(firResolveState)
     private val firScopeStorage = FirScopeRegistry()
 
-    private val memberScopeCache = IdentityHashMap<KtClassOrObjectSymbol, KtMemberScope>()
+    private val memberScopeCache = IdentityHashMap<KtSymbolWithDeclarations, KtMemberScope>()
-    private val declaredMemberScopeCache = IdentityHashMap<KtClassOrObjectSymbol, KtDeclaredMemberScope>()
+    private val declaredMemberScopeCache = IdentityHashMap<KtSymbolWithDeclarations, KtDeclaredMemberScope>()
     private val packageMemberScopeCache = IdentityHashMap<KtPackageSymbol, KtPackageScope>()
 
-    override fun getMemberScope(classSymbol: KtClassOrObjectSymbol): KtMemberScope = withValidityAssertion {
+    private fun KtSymbolWithDeclarations.firRef(): FirRefWithValidityCheck<out FirClass<*>>? = when (this) {
+        is KtFirClassOrObjectSymbol -> this.firRef
+        is KtFirEnumEntrySymbol -> this.initializerFirRef
+        else -> error { "Unknown KtSymbolWithDeclarations implementation ${this::class.qualifiedName}" }
+    }
+
+    override fun getMemberScope(classSymbol: KtSymbolWithDeclarations): KtMemberScope = withValidityAssertion {
         memberScopeCache.getOrPut(classSymbol) {
-            check(classSymbol is KtFirClassOrObjectSymbol)
+
+            val firRef = classSymbol.firRef()
+                ?: return@getOrPut KtFirEmptyMemberScope(classSymbol)
+
             val firScope =
-                classSymbol.firRef.withFir(FirResolvePhase.SUPER_TYPES) { fir ->
+                firRef.withFir(FirResolvePhase.SUPER_TYPES) { fir ->
                     val firSession = fir.session
                     fir.unsubstitutedScope(
                         firSession,
@@ -70,10 +82,12 @@ internal class KtFirScopeProvider(
         }
     }
 
-    override fun getDeclaredMemberScope(classSymbol: KtClassOrObjectSymbol): KtDeclaredMemberScope = withValidityAssertion {
+    override fun getDeclaredMemberScope(classSymbol: KtSymbolWithDeclarations): KtDeclaredMemberScope = withValidityAssertion {
         declaredMemberScopeCache.getOrPut(classSymbol) {
-            check(classSymbol is KtFirClassOrObjectSymbol)
+            val firRef = classSymbol.firRef()
-            val firScope = classSymbol.firRef.withFir(FirResolvePhase.SUPER_TYPES) { declaredMemberScope(it) }
+                ?: return@getOrPut KtFirEmptyMemberScope(classSymbol)
+
+            val firScope = firRef.withFir(FirResolvePhase.SUPER_TYPES) { declaredMemberScope(it) }
                 .also(firScopeStorage::register)
             KtFirDeclaredMemberScope(classSymbol, firScope, token, builder)
         }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirDeclaredMemberScope.kt

@@ -13,9 +13,10 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymb
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtDeclaredMemberScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtUnsubstitutedScope
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 
 internal class KtFirDeclaredMemberScope(
-    override val owner: KtFirClassOrObjectSymbol,
+    override val owner: KtSymbolWithDeclarations,
     firScope: FirClassDeclaredMemberScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirEmptyMemberScope.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2010-2020 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.idea.frontend.api.fir.scopes
+
+import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.ValidityTokenOwner
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtDeclaredMemberScope
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtMemberScope
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassifierSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
+import org.jetbrains.kotlin.name.Name
+
+internal class KtFirEmptyMemberScope(override val owner: KtSymbolWithDeclarations) : KtMemberScope, KtDeclaredMemberScope, ValidityTokenOwner {
+    override fun getCallableNames(): Set<Name> = emptySet()
+
+    override fun getClassifierNames(): Set<Name> = emptySet()
+
+    override fun getCallableSymbols(nameFilter: KtScopeNameFilter): Sequence<KtCallableSymbol> =
+        emptySequence()
+
+    override fun getClassifierSymbols(nameFilter: KtScopeNameFilter): Sequence<KtClassifierSymbol> =
+        emptySequence()
+
+    override val token: ValidityToken
+        get() = owner.token
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/scopes/KtFirMemberScope.kt

@@ -13,9 +13,10 @@ import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymb
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.weakRef
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtMemberScope
 import org.jetbrains.kotlin.idea.frontend.api.scopes.KtUnsubstitutedScope
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithDeclarations
 
 internal class KtFirMemberScope(
-    override val owner: KtFirClassOrObjectSymbol,
+    override val owner: KtSymbolWithDeclarations,
     firScope: FirTypeScope,
     token: ValidityToken,
     builder: KtSymbolByFirBuilder

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/KtFirEnumEntrySymbol.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.frontend.api.fir.symbols
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.containingClass
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousObject
 import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.findPsi
@@ -15,6 +16,7 @@ import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.KtFirEnumEntrySymbolPointer
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.pointers.createSignature
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.FirRefWithValidityCheck
 import org.jetbrains.kotlin.idea.frontend.api.fir.utils.firRef
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtEnumEntrySymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.pointers.KtPsiBasedSymbolPointer
@@ -30,6 +32,7 @@ internal class KtFirEnumEntrySymbol(
     private val builder: KtSymbolByFirBuilder
 ) : KtEnumEntrySymbol(), KtFirSymbol<FirEnumEntry> {
     override val firRef = firRef(fir, resolveState)
+
     override val psi: PsiElement? by firRef.withFirAndCache { fir -> fir.findPsi(fir.session) }
 
     override val name: Name get() = firRef.withFir { it.name }
@@ -38,6 +41,8 @@ internal class KtFirEnumEntrySymbol(
     override val containingEnumClassIdIfNonLocal: ClassId?
         get() = firRef.withFir { it.containingClass()?.classId?.takeUnless { it.isLocal } }
 
+    override val hasBody: Boolean get() = firRef.withFir { it.initializer != null }
+
     override fun createPointer(): KtSymbolPointer<KtEnumEntrySymbol> {
         KtPsiBasedSymbolPointer.createForSymbolFromSource(this)?.let { return it }
         return KtFirEnumEntrySymbolPointer(containingEnumClassIdIfNonLocal!!, firRef.withFir { it.createSignature() })

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/symbols/firSymbolUtils.kt

@@ -22,7 +22,7 @@ internal inline fun <reified M : KtSymbolModality> Modality?.getSymbolModality()
     Modality.OPEN -> KtCommonSymbolModality.OPEN
     Modality.ABSTRACT -> KtCommonSymbolModality.ABSTRACT
     Modality.SEALED -> KtSymbolModality.SEALED
-    null -> error("Symbol modality should not be null, looks like the fir symbol was not properly resolved")
+    null -> KtCommonSymbolModality.UNKNOWN
 } as? M ?: error("Sealed modality can only be applied to class")
 
 internal inline fun <F : FirMemberDeclaration, reified M : KtSymbolModality> KtFirSymbol<F>.getModality() =

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/utils/firUtils.kt

@@ -5,23 +5,16 @@
 package org.jetbrains.kotlin.idea.frontend.api.fir.utils
 
 import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.FirSymbolOwner
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirAnnotatedDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.getPrimaryConstructorIfAny
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.resolve.toSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtSymbolByFirBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirAnnotationCall
-import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedConstantValue
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.*
-import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSimpleConstantValue
-import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtUnsupportedConstantValue
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtCallElement
 
 internal fun mapAnnotationParameters(annotationCall: FirAnnotationCall, session: FirSession): Map<String, FirExpression> {
@@ -54,15 +47,20 @@ internal fun mapAnnotationParameters(annotationCall: FirAnnotationCall, session:
     return resultSet
 }
 
-private fun convertAnnotation(annotationCall: FirAnnotationCall, session: FirSession): KtFirAnnotationCall? {
+private fun FirExpression.convertConstantExpression(): KtConstantValue =
+    when (this) {
+        is FirConstExpression<*> -> KtSimpleConstantValue(value)
+        else -> KtUnsupportedConstantValue
+    }
+
+private fun convertAnnotation(
+    annotationCall: FirAnnotationCall,
+    session: FirSession
+): KtFirAnnotationCall? {
 
     val annotationCone = annotationCall.annotationTypeRef.coneType as? ConeClassLikeType ?: return null
     val classId = annotationCone.classId ?: return null
 
-    fun FirExpression.convertConstantExpression() =
-        (this as? FirConstExpression<*>)?.value?.let { KtSimpleConstantValue(it) }
-            ?: KtUnsupportedConstantValue
-
     val resultList = mapAnnotationParameters(annotationCall, session).map {
         KtNamedConstantValue(it.key, it.value.convertConstantExpression())
     }
@@ -75,6 +73,7 @@ private fun convertAnnotation(annotationCall: FirAnnotationCall, session: FirSes
     )
 }
 
-internal fun convertAnnotation(declaration: FirAnnotatedDeclaration) = declaration.annotations.mapNotNull {
+internal fun convertAnnotation(declaration: FirAnnotatedDeclaration): List<KtFirAnnotationCall> =
-    convertAnnotation(it, declaration.session)
+    declaration.annotations.mapNotNull {
-}
+        convertAnnotation(it, declaration.session)
+    }