FIR IDE: implement completion in sealed when

2021-05-27 15:06:15 +02:00
parent 4bdf46e0a4
commit 9af0885e82
11 changed files with 407 additions and 15 deletions
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2021 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.completion
+
+import com.intellij.codeInsight.completion.InsertHandler
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.lookup.LookupElement
+import org.jetbrains.kotlin.idea.completion.lookups.KotlinLookupObject
+import org.jetbrains.kotlin.psi.KtFile
+import kotlin.reflect.KClass
+
+internal abstract class InsertionHandlerBase<LO : KotlinLookupObject>(
+    private val lookupObjectClass: KClass<LO>
+) : InsertHandler<LookupElement> {
+
+    final override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        val ktFile = context.file as? KtFile ?: return
+        val lookupObject = item.`object`
+        check(lookupObjectClass.isInstance(lookupObject))
+
+        @Suppress("UNCHECKED_CAST")
+        handleInsert(context, item, ktFile, lookupObject as LO)
+    }
+
+    abstract fun handleInsert(context: InsertionContext, item: LookupElement, ktFile: KtFile, lookupObject: LO)
+}
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.originalKtFile
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtFile
-import java.io.File

 class KotlinFirCompletionContributor : CompletionContributor() {
    init {
@@ -71,6 +70,7 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
        val importDirectivePackageMembersCompletionContributor = FirImportDirectivePackageMembersCompletionContributor(basicContext)
        val typeParameterConstraintNameInWhereClauseContributor = FirTypeParameterConstraintNameInWhereClauseCompletionContributor(basicContext)
        val classifierNameContributor = FirSameAsFileClassifierNameCompletionContributor(basicContext)
+        val whenWithSubjecConditionContributor = FirWhenWithSubjectConditionContributor(basicContext)

        when (positionContext) {
            is FirExpressionNameReferencePositionContext -> {
@@ -117,6 +117,10 @@ private object KotlinFirCompletionProvider : CompletionProvider<CompletionParame
                complete(classifierNameContributor, positionContext)
            }

+            is FirWithSubjectEntryPositionContext -> {
+                complete(whenWithSubjecConditionContributor, positionContext)
+            }
+
            is FirIncorrectPositionContext -> {
                // do nothing, completion is not suposed to be called here
            }
@@ -88,6 +88,14 @@ internal class FirExpressionNameReferencePositionContext(
 ) : FirNameReferencePositionContext()


+internal class FirWithSubjectEntryPositionContext(
+    override val position: PsiElement,
+    override val reference: KtSimpleNameReference,
+    override val nameExpression: KtSimpleNameExpression,
+    override val explicitReceiver: KtExpression?,
+    val whenCondition: KtWhenCondition,
+) : FirNameReferencePositionContext()
+
 internal class FirUnknownPositionContext(
    override val position: PsiElement
 ) : FirRawPositionCompletionContext()
@@ -122,6 +130,14 @@ internal object FirPositionCompletionContextDetector {
            parent is KtUserType -> {
                detectForTypeContext(parent, position, reference, nameExpression, explicitReceiver)
            }
+            parent is KtWhenCondition && parent.isConditionOnWhenWithSubject() -> {
+                FirWithSubjectEntryPositionContext(
+                    position,
+                    reference,
+                    nameExpression,
+                    explicitReceiver, parent
+                )
+            }
            nameExpression.isReferenceExpressionInImportDirective() -> {
                FirImportDirectivePositionContext(
                    position,
@@ -149,6 +165,12 @@ internal object FirPositionCompletionContextDetector {
        }
    }

+    private fun KtWhenCondition.isConditionOnWhenWithSubject(): Boolean {
+        val whenEntry = (parent as? KtWhenEntry) ?: return false
+        val whenExpression = whenEntry.parent as? KtWhenExpression ?: return false
+        return whenExpression.subjectExpression != null
+    }
+
    private fun KtExpression.isReferenceExpressionInImportDirective() = when (val parent = parent) {
        is KtImportDirective -> parent.importedReference == this
        is KtDotQualifiedExpression -> {
@@ -12,8 +12,8 @@ import org.jetbrains.kotlin.idea.completion.contributors.helpers.getStaticScope
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
 import org.jetbrains.kotlin.idea.frontend.api.types.KtClassType
-import org.jetbrains.kotlin.psi.KtEnumEntry
 import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.FirClassifierProvider.getAvailableClassifiersCurrentScope

 internal open class FirClassifierCompletionContributor(
    basicContext: FirBasicCompletionContext,
@@ -51,19 +51,14 @@ internal open class FirClassifierCompletionContributor(
        positionContext: FirNameReferencePositionContext,
        visibilityChecker: CompletionVisibilityChecker
    ) {
-        val implicitScopes = originalKtFile.getScopeContextForPosition(positionContext.nameExpression).scopes
-        val classesFromScopes = implicitScopes
-            .getClassifierSymbols(scopeNameFilter)
+        getAvailableClassifiersCurrentScope(
+            originalKtFile,
+            positionContext.nameExpression,
+            scopeNameFilter,
+            indexHelper,
+            visibilityChecker
+        )
            .filter { filterClassifiers(it) }
-            .filter { with(visibilityChecker) { isVisible(it) } }
-
-        classesFromScopes.forEach { addClassifierSymbolToCompletion(it, insertFqName = true) }
-
-        val kotlinClassesFromIndices = indexHelper.getKotlinClasses(scopeNameFilter, psiFilter = { it !is KtEnumEntry })
-        kotlinClassesFromIndices.asSequence()
-            .map { it.getSymbol() as KtClassifierSymbol }
-            .filter { filterClassifiers(it) }
-            .filter { with(visibilityChecker) { isVisible(it) } }
            .forEach { addClassifierSymbolToCompletion(it, insertFqName = true) }
    }
 }
@@ -0,0 +1,230 @@
+/*
+ * Copyright 2010-2021 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.completion.contributors
+
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.openapi.util.TextRange
+import com.intellij.psi.util.parentOfType
+import gnu.trove.THashSet
+import gnu.trove.TObjectHashingStrategy
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.idea.completion.InsertionHandlerBase
+import org.jetbrains.kotlin.idea.completion.KotlinFirIconProvider.getIconFor
+import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker
+import org.jetbrains.kotlin.idea.completion.context.FirBasicCompletionContext
+import org.jetbrains.kotlin.idea.completion.context.FirWithSubjectEntryPositionContext
+import org.jetbrains.kotlin.idea.completion.createKeywordElement
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.*
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.idea.completion.contributors.helpers.FirClassifierProvider.getAvailableClassifiersCurrentScope
+import org.jetbrains.kotlin.idea.completion.lookups.KotlinLookupObject
+import org.jetbrains.kotlin.idea.completion.lookups.shortenReferencesForFirCompletion
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
+import org.jetbrains.kotlin.idea.frontend.api.types.*
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.renderer.render
+
+internal class FirWhenWithSubjectConditionContributor(
+    basicContext: FirBasicCompletionContext,
+) : FirContextCompletionContributorBase<FirWithSubjectEntryPositionContext>(basicContext) {
+    override fun KtAnalysisSession.complete(positionContext: FirWithSubjectEntryPositionContext) {
+        val whenCondition = positionContext.whenCondition
+        val whenExpression = whenCondition.parentOfType<KtWhenExpression>() ?: return
+        val subject = whenExpression.subjectExpression ?: return
+        val allConditionsExceptCurrent = whenExpression.entries.flatMap { entry -> entry.conditions.filter { it != whenCondition } }
+        val subjectType = subject.getKtType()
+        val classSymbol = getClassSymbol(subjectType)
+        val visibilityChecker = CompletionVisibilityChecker.create(basicContext, positionContext)
+        when {
+            classSymbol?.classKind == KtClassKind.ENUM_CLASS -> {
+                completeEnumEntries(classSymbol, allConditionsExceptCurrent, visibilityChecker)
+            }
+            classSymbol?.modality == Modality.SEALED -> {
+                completeSubClassesOfSealedClass(classSymbol, allConditionsExceptCurrent, whenCondition, visibilityChecker)
+            }
+            else -> {
+                completeAllTypes(whenCondition, visibilityChecker)
+            }
+        }
+        addNullIfWhenExpressionCanReturnNull(subjectType)
+        addElseBranchIfSingleConditionInEntry(whenCondition)
+    }
+
+    private fun KtAnalysisSession.getClassSymbol(subjectType: KtType): KtNamedClassOrObjectSymbol? {
+        val classType = subjectType as? KtClassType
+        return classType?.classSymbol as? KtNamedClassOrObjectSymbol
+    }
+
+
+    private fun KtAnalysisSession.addNullIfWhenExpressionCanReturnNull(type: KtType?) {
+        if (type?.canBeNull == true) {
+            val lookupElement = createKeywordElement(keyword = KtTokens.NULL_KEYWORD.value)
+            result.addElement(lookupElement)
+        }
+    }
+
+    private fun KtAnalysisSession.completeAllTypes(
+        whenCondition: KtWhenCondition,
+        visibilityChecker: CompletionVisibilityChecker,
+    ) {
+        getAvailableClassifiersCurrentScope(originalKtFile, whenCondition, scopeNameFilter, indexHelper, visibilityChecker)
+            .forEach { classifier ->
+                if (classifier !is KtNamedSymbol) return@forEach
+
+                addLookupElement(
+                    classifier.name.asString(),
+                    classifier,
+                    (classifier as? KtNamedClassOrObjectSymbol)?.classIdIfNonLocal?.asSingleFqName(),
+                    isPrefixNeeded(classifier)
+                )
+            }
+    }
+
+    private fun KtAnalysisSession.isPrefixNeeded(classifier: KtClassifierSymbol): Boolean {
+        return when (classifier) {
+            is KtAnonymousObjectSymbol -> return false
+            is KtNamedClassOrObjectSymbol -> !classifier.classKind.isObject
+            is KtTypeAliasSymbol -> (classifier.expandedType as? KtClassType)?.classSymbol?.let { isPrefixNeeded(it) } == true
+            is KtTypeParameterSymbol -> true
+        }
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun KtAnalysisSession.completeSubClassesOfSealedClass(
+        classSymbol: KtNamedClassOrObjectSymbol,
+        conditions: List<KtWhenCondition>,
+        whenCondition: KtWhenCondition,
+        visibilityChecker: CompletionVisibilityChecker,
+    ) {
+        require(classSymbol.modality == Modality.SEALED)
+        val handledCasesClassIds = getHandledClassIds(conditions)
+        val allInheritors = getAllSealedInheritors(classSymbol)
+
+        allInheritors
+            .asSequence()
+            .filter { it.classIdIfNonLocal !in handledCasesClassIds }
+            .filter { with(visibilityChecker) { isVisible(it as KtClassifierSymbol) } }
+            .forEach { inheritor ->
+                val classId = inheritor.classIdIfNonLocal ?: return@forEach
+                addLookupElement(classId.relativeClassName.asString(), inheritor, classId.asSingleFqName(), isPrefixNeeded(inheritor))
+            }
+
+        if (allInheritors.any { it.modality == Modality.ABSTRACT }) {
+            completeAllTypes(whenCondition, visibilityChecker)
+        }
+    }
+
+    private fun KtAnalysisSession.getHandledClassIds(conditions: List<KtWhenCondition>): Set<ClassId> =
+        conditions.mapNotNullTo(hashSetOf()) { condition ->
+            val reference = when (condition) {
+                is KtWhenConditionWithExpression -> condition.expression?.reference()
+                is KtWhenConditionIsPattern -> (condition.typeReference?.typeElement as? KtUserType)?.referenceExpression?.reference()
+                else -> null
+            }
+            val resolvesTo = reference?.resolveToSymbol() as? KtNamedClassOrObjectSymbol
+            resolvesTo?.classIdIfNonLocal
+        }
+
+    private fun KtAnalysisSession.getAllSealedInheritors(classSymbol: KtNamedClassOrObjectSymbol): Collection<KtNamedClassOrObjectSymbol> {
+        fun KtAnalysisSession.getAllSealedInheritorsTo(
+            classSymbol: KtNamedClassOrObjectSymbol,
+            destination: MutableSet<KtNamedClassOrObjectSymbol>
+        ) {
+            classSymbol.getSealedClassInheritors().forEach { inheritor ->
+                destination += inheritor
+                if (inheritor.modality == Modality.SEALED) {
+                    getAllSealedInheritorsTo(inheritor, destination)
+                }
+            }
+        }
+
+        return THashSet(KtNamedClassOrObjectSymbolTObjectHashingStrategy)
+            .apply { getAllSealedInheritorsTo(classSymbol, this) }
+    }
+
+    private fun addElseBranchIfSingleConditionInEntry(whenCondition: KtWhenCondition) {
+        val whenEntry = whenCondition.parent as? KtWhenEntry ?: return
+        if (whenEntry.conditions.size > 1) return
+        val lookupElement = createKeywordElement(keyword = KtTokens.ELSE_KEYWORD.value, tail = " -> ")
+        result.addElement(lookupElement)
+    }
+
+
+    private fun KtAnalysisSession.completeEnumEntries(
+        classSymbol: KtNamedClassOrObjectSymbol,
+        conditions: List<KtWhenCondition>,
+        visibilityChecker: CompletionVisibilityChecker
+    ) {
+        require(classSymbol.classKind == KtClassKind.ENUM_CLASS)
+        val handledCasesNames = conditions.mapNotNullTo(hashSetOf()) { condition ->
+            val conditionWithExpression = condition as? KtWhenConditionWithExpression
+            val resolvesTo = conditionWithExpression?.expression?.reference()?.resolveToSymbol() as? KtEnumEntrySymbol
+            resolvesTo?.name
+        }
+        val allEnumEntrySymbols = classSymbol.getEnumEntries()
+        allEnumEntrySymbols
+            .filter { it.name !in handledCasesNames }
+            .filter { with(visibilityChecker) { isVisible(it) } }
+            .forEach { entry ->
+                addLookupElement(
+                    "${classSymbol.name.asString()}.${entry.name.asString()}",
+                    entry,
+                    entry.callableIdIfNonLocal?.asSingleFqName(),
+                    isPrefixNeeded = false
+                )
+            }
+    }
+
+    private fun KtAnalysisSession.addLookupElement(lookupString: String, symbol: KtNamedSymbol, fqName: FqName?, isPrefixNeeded: Boolean) {
+        val lookupObject = WhenConditionLookupObject(symbol.name, fqName, isPrefixNeeded)
+
+        LookupElementBuilder.create(lookupObject, getIsPrefix(isPrefixNeeded) + lookupString)
+            .withIcon(getIconFor(symbol))
+            .withPsiElement(symbol.psi)
+            .withInsertHandler(WhenConditionInsertionHandler)
+            .let(result::addElement)
+    }
+}
+
+private data class WhenConditionLookupObject(
+    override val shortName: Name,
+    val fqName: FqName?,
+    val needIsPrefix: Boolean,
+) : KotlinLookupObject
+
+private object WhenConditionInsertionHandler : InsertionHandlerBase<WhenConditionLookupObject>(WhenConditionLookupObject::class) {
+    override fun handleInsert(context: InsertionContext, item: LookupElement, ktFile: KtFile, lookupObject: WhenConditionLookupObject) {
+        if (lookupObject.fqName != null) {
+            val fqName = lookupObject.fqName
+            context.document.replaceString(
+                context.startOffset,
+                context.tailOffset,
+                getIsPrefix(lookupObject.needIsPrefix) + fqName.render()
+            )
+            context.commitDocument()
+
+            shortenReferencesForFirCompletion(ktFile, TextRange(context.startOffset, context.tailOffset))
+        }
+    }
+}
+
+private fun getIsPrefix(prefixNeeded: Boolean): String {
+    return if (prefixNeeded) "is " else ""
+}
+
+private object KtNamedClassOrObjectSymbolTObjectHashingStrategy : TObjectHashingStrategy<KtNamedClassOrObjectSymbol> {
+    override fun equals(p0: KtNamedClassOrObjectSymbol, p1: KtNamedClassOrObjectSymbol): Boolean =
+        p0.classIdIfNonLocal == p1.classIdIfNonLocal
+
+    override fun computeHashCode(p0: KtNamedClassOrObjectSymbol): Int =
+        p0.classIdIfNonLocal.hashCode()
+}
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2021 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.completion.contributors.helpers
+
+import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.scopes.KtScopeNameFilter
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassifierSymbol
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtEnumEntry
+import org.jetbrains.kotlin.psi.KtFile
+
+internal object FirClassifierProvider {
+    fun KtAnalysisSession.getAvailableClassifiersCurrentScope(
+        originalKtFile: KtFile,
+        position: KtElement,
+        scopeNameFilter: KtScopeNameFilter,
+        indexHelper: IndexHelper,
+        visibilityChecker: CompletionVisibilityChecker
+    ): Sequence<KtClassifierSymbol> = sequence {
+        yieldAll(
+            originalKtFile.getScopeContextForPosition(position).scopes
+                .getClassifierSymbols(scopeNameFilter)
+                .filter { with(visibilityChecker) { isVisible(it) } }
+        )
+
+        yieldAll(
+            indexHelper.getKotlinClasses(scopeNameFilter, psiFilter = { it !is KtEnumEntry }).asSequence()
+                .map { it.getSymbol() as KtClassifierSymbol }
+                .filter { with(visibilityChecker) { isVisible(it) } }
+        )
+    }
+}
@@ -42,7 +42,8 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
    KtReferenceShortenerMixIn,
    KtSymbolDeclarationRendererMixIn,
    KtVisibilityCheckerMixIn,
-    KtMemberSymbolProviderMixin
+    KtMemberSymbolProviderMixin,
+    KtInheritorsProviderMixIn
 {

    override val analysisSession: KtAnalysisSession get() = this
@@ -99,4 +100,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali

    internal val overrideInfoProvider: KtOverrideInfoProvider get() = overrideInfoProviderImpl
    protected abstract val overrideInfoProviderImpl: KtOverrideInfoProvider
+
+    internal val inheritorsProvider: KtInheritorsProvider get() = inheritorsProviderImpl
+    protected abstract val inheritorsProviderImpl: KtInheritorsProvider
 }
@@ -0,0 +1,22 @@
+/*
+ * 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.components
+
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtNamedClassOrObjectSymbol
+
+abstract class KtInheritorsProvider : KtAnalysisSessionComponent() {
+    abstract fun getInheritorsOfSealedClass(classSymbol: KtNamedClassOrObjectSymbol): List<KtNamedClassOrObjectSymbol>
+    abstract fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol>
+}
+
+interface KtInheritorsProviderMixIn : KtAnalysisSessionMixIn {
+     fun KtNamedClassOrObjectSymbol.getSealedClassInheritors(): List<KtNamedClassOrObjectSymbol> =
+         analysisSession.inheritorsProvider.getInheritorsOfSealedClass(this)
+
+     fun KtNamedClassOrObjectSymbol.getEnumEntries(): List<KtEnumEntrySymbol> =
+         analysisSession.inheritorsProvider.getEnumEntries(this)
+}
@@ -34,6 +34,7 @@ class SealedClassInheritorsProviderIdeImpl : SealedClassInheritorsProvider() {
    }

    private fun getInheritors(firClass: FirRegularClass): List<ClassId> {
+        // TODO fix for non-source classes
        val sealedKtClass = firClass.psi as? KtClass ?: return emptyList()
        val module = sealedKtClass.module ?: return emptyList()
        val containingPackage = firClass.classId.packageFqName
@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.*
 import org.jetbrains.kotlin.idea.frontend.api.InvalidWayOfUsingAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.components.KtInheritorsProvider
 import org.jetbrains.kotlin.idea.frontend.api.components.KtVisibilityChecker
 import org.jetbrains.kotlin.idea.frontend.api.components.KtSymbolDeclarationRendererProvider
 import org.jetbrains.kotlin.idea.frontend.api.fir.components.*
@@ -78,6 +79,8 @@ private constructor(

    override val subtypingComponentImpl = KtFirSubtypingComponent(this, token)

+    override val inheritorsProviderImpl: KtInheritorsProvider = KtFirInheritorsProvider(this, token)
+
    override fun createContextDependentCopy(originalKtFile: KtFile, elementToReanalyze: KtElement): KtAnalysisSession {
        check(mode == AnalysisSessionMode.REGULAR) {
            "Cannot create context-dependent copy of KtAnalysis session from a context dependent one"
@@ -0,0 +1,45 @@
+/*
+ * 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.components
+
+import org.jetbrains.kotlin.descriptors.Modality
+import org.jetbrains.kotlin.fir.declarations.getSealedClassInheritors
+import org.jetbrains.kotlin.idea.frontend.api.components.KtInheritorsProvider
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassKind
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtEnumEntrySymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.withValidityAssertion
+
+internal class KtFirInheritorsProvider(
+    override val analysisSession: KtFirAnalysisSession,
+    override val token: ValidityToken,
+) : KtInheritorsProvider(), KtFirAnalysisSessionComponent {
+    override fun getInheritorsOfSealedClass(
+        classSymbol: KtNamedClassOrObjectSymbol
+    ): List<KtNamedClassOrObjectSymbol> = withValidityAssertion {
+        require(classSymbol.modality == Modality.SEALED)
+        require(classSymbol is KtFirNamedClassOrObjectSymbol)
+
+        val inheritorClassIds = classSymbol.firRef.withFir { fir ->
+            fir.getSealedClassInheritors(analysisSession.rootModuleSession)
+        }
+
+        with(analysisSession) {
+            inheritorClassIds.mapNotNull { it.getCorrespondingToplevelClassOrObjectSymbol() as? KtNamedClassOrObjectSymbol }
+        }
+    }
+
+    override fun getEnumEntries(classSymbol: KtNamedClassOrObjectSymbol): List<KtEnumEntrySymbol> = withValidityAssertion {
+        require(classSymbol.classKind == KtClassKind.ENUM_CLASS)
+        with(analysisSession) {
+            classSymbol.getDeclaredMemberScope().getCallableSymbols().filterIsInstance<KtEnumEntrySymbol>().toList()
+        }
+    }
+
+}