[AA Standalone] Implement KotlinDirectInheritorsProvider

- We are relying on static indexing to find candidates for sealed
  inheritors, hence the extension to the index.
- The direct usage of `KotlinStaticDeclarationProviderFactory` in
  `KotlinStandaloneDirectInheritorsProvider` is not pretty, but a proper
  design requires making the static index available as a service and
  moving "static" services to the Standalone API (from AA providers).

^KT-66013

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirSymbolDeclarationOverridesProvider.kt

@@ -19,13 +19,12 @@ import org.jetbrains.kotlin.analysis.api.symbols.KtSymbolOrigin
 import org.jetbrains.kotlin.analysis.api.symbols.KtValueParameterSymbol
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.utils.superConeTypes
 import org.jetbrains.kotlin.fir.scopes.*
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
-import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.unwrapFakeOverrides
+import org.jetbrains.kotlin.analysis.api.fir.utils.isSubClassOf
 
 internal class KtFirSymbolDeclarationOverridesProvider(
     override val analysisSession: KtFirAnalysisSession,
@@ -168,37 +167,23 @@ internal class KtFirSymbolDeclarationOverridesProvider(
     }
 
     override fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
-        return isSubClassOf(subClass, superClass, checkDeep = true)
+        return isSubClassOf(subClass, superClass, allowIndirectSubtyping = true)
     }
 
     override fun isDirectSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol): Boolean {
-        return isSubClassOf(subClass, superClass, checkDeep = false)
+        return isSubClassOf(subClass, superClass, allowIndirectSubtyping = false)
     }
 
-    private fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol, checkDeep: Boolean): Boolean {
+    private fun isSubClassOf(subClass: KtClassOrObjectSymbol, superClass: KtClassOrObjectSymbol, allowIndirectSubtyping: Boolean): Boolean {
         require(subClass is KtFirSymbol<*>)
         require(superClass is KtFirSymbol<*>)
 
         if (subClass == superClass) return false
-        subClass.firSymbol.lazyResolveToPhase(FirResolvePhase.SUPER_TYPES)
         return isSubClassOf(
             subClass = subClass.firSymbol.fir as FirClass,
             superClass = superClass.firSymbol.fir as FirClass,
-            checkDeep
+            allowIndirectSubtyping,
         )
-
-
-    }
-
-    private fun isSubClassOf(subClass: FirClass, superClass: FirClass, checkDeep: Boolean): Boolean {
-        if (subClass.superConeTypes.any { it.toRegularClassSymbol(rootModuleSession) == superClass.symbol }) return true
-        if (!checkDeep) return false
-        subClass.superConeTypes.forEach { superType ->
-            val superOfSub = superType.toRegularClassSymbol(rootModuleSession) ?: return@forEach
-            superOfSub.lazyResolveToPhase(FirResolvePhase.SUPER_TYPES)
-            if (isSubClassOf(superOfSub.fir, superClass, checkDeep = true)) return true
-        }
-        return false
     }
 
     override fun getIntersectionOverriddenSymbols(symbol: KtCallableSymbol): List<KtCallableSymbol> {

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/utils/typeUtils.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.fir.utils
+
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.declarations.utils.superConeTypes
+import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
+import org.jetbrains.kotlin.fir.types.toRegularClassSymbol
+
+/**
+ * Returns whether [subClass] is a strict subtype of [superClass]. Resolves [subClass] to [FirResolvePhase.SUPER_TYPES].
+ */
+fun isSubClassOf(subClass: FirClass, superClass: FirClass, allowIndirectSubtyping: Boolean = true): Boolean {
+    subClass.lazyResolveToPhase(FirResolvePhase.SUPER_TYPES)
+
+    val session = subClass.moduleData.session
+    if (subClass.superConeTypes.any { it.toRegularClassSymbol(session) == superClass.symbol }) return true
+    if (!allowIndirectSubtyping) return false
+
+    subClass.superConeTypes.forEach { superType ->
+        val superOfSub = superType.toRegularClassSymbol(session) ?: return@forEach
+        if (isSubClassOf(superOfSub.fir, superClass, allowIndirectSubtyping = true)) return true
+    }
+    return false
+}

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationIndex.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.analysis.providers.impl
 
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
 internal class KotlinStaticDeclarationIndex {
@@ -16,4 +17,9 @@ internal class KotlinStaticDeclarationIndex {
     internal val typeAliasMap: MutableMap<FqName, MutableSet<KtTypeAlias>> = mutableMapOf()
     internal val topLevelFunctionMap: MutableMap<FqName, MutableSet<KtNamedFunction>> = mutableMapOf()
     internal val topLevelPropertyMap: MutableMap<FqName, MutableSet<KtProperty>> = mutableMapOf()
+
+    /**
+     * Allows quickly finding [KtClassOrObject]s which have a given simple name as a supertype. The map may contain local classes as well.
+     */
+    internal val classesBySupertypeName: MutableMap<Name, MutableSet<KtClassOrObject>> = mutableMapOf()
 }

analysis/analysis-api-providers/src/org/jetbrains/kotlin/analysis/providers/impl/KotlinStaticDeclarationProvider.kt

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.name.*
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getSuperNames
 import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.impl.*
@@ -203,7 +204,7 @@ public class KotlinStaticDeclarationProviderFactory(
         }
 
         override fun visitClassOrObject(classOrObject: KtClassOrObject) {
-            addToClassMap(classOrObject)
+            indexClassOrObject(classOrObject)
             super.visitClassOrObject(classOrObject)
         }
 
@@ -236,6 +237,11 @@ public class KotlinStaticDeclarationProviderFactory(
         }.add(script)
     }
 
+    private fun indexClassOrObject(classOrObject: KtClassOrObject) {
+        addToClassMap(classOrObject)
+        indexSupertypeNames(classOrObject)
+    }
+
     private fun addToClassMap(classOrObject: KtClassOrObject) {
         classOrObject.getClassId()?.let { classId ->
             index.classMap.computeIfAbsent(classId.packageFqName) {
@@ -244,6 +250,14 @@ public class KotlinStaticDeclarationProviderFactory(
         }
     }
 
+    private fun indexSupertypeNames(classOrObject: KtClassOrObject) {
+        classOrObject.getSuperNames().forEach { superName ->
+            index.classesBySupertypeName
+                .computeIfAbsent(Name.identifier(superName)) { mutableSetOf() }
+                .add(classOrObject)
+        }
+    }
+
     private fun addToTypeAliasMap(typeAlias: KtTypeAlias) {
         typeAlias.getClassId()?.let { classId ->
             index.typeAliasMap.computeIfAbsent(classId.packageFqName) {
@@ -307,12 +321,12 @@ public class KotlinStaticDeclarationProviderFactory(
     private fun indexStub(stub: StubElement<*>) {
         when (stub) {
             is KotlinClassStubImpl -> {
-                addToClassMap(stub.psi)
+                indexClassOrObject(stub.psi)
                 // member functions and properties
                 stub.childrenStubs.forEach(::indexStub)
             }
             is KotlinObjectStubImpl -> {
-                addToClassMap(stub.psi)
+                indexClassOrObject(stub.psi)
                 // member functions and properties
                 stub.childrenStubs.forEach(::indexStub)
             }
@@ -399,6 +413,9 @@ public class KotlinStaticDeclarationProviderFactory(
     }
 
     public fun getAllKtClasses(): List<KtClassOrObject> = index.classMap.values.flattenTo(mutableListOf())
+
+    public fun getDirectInheritorCandidates(baseClassName: Name): Set<KtClassOrObject> =
+        index.classesBySupertypeName[baseClassName].orEmpty()
 }
 
 /**

analysis/analysis-api-standalone/analysis-api-fir-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/FirStandaloneServiceRegistrar.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.api.fir.references.ReadWriteAccessCheckerFirImpl
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.analysis.api.standalone.base.providers.KotlinStandaloneDirectInheritorsProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService
@@ -24,6 +25,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCach
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.stubBased.deserialization.LLStubBasedLibrarySymbolProviderFactory
+import org.jetbrains.kotlin.analysis.providers.KotlinDirectInheritorsProvider
 import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
@@ -63,6 +65,8 @@ object FirStandaloneServiceRegistrar : AnalysisApiStandaloneServiceRegistrar {
             LLFirSessionInvalidationService.getInstance(project).subscribeToModificationEvents()
 
             registerService(LLFirDeclarationModificationService::class.java)
+
+            registerService(KotlinDirectInheritorsProvider::class.java, KotlinStandaloneDirectInheritorsProvider(project))
         }
     }
 

analysis/analysis-api-standalone/analysis-api-fir-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/providers/KotlinStandaloneDirectInheritorsProvider.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2010-2024 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.analysis.api.standalone.base.providers
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.analysis.api.fir.utils.isSubClassOf
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
+import org.jetbrains.kotlin.analysis.project.structure.KtDanglingFileModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProviderFactory
+import org.jetbrains.kotlin.analysis.providers.KotlinDirectInheritorsProvider
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticDeclarationProviderFactory
+import org.jetbrains.kotlin.fir.declarations.FirClass
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.SymbolInternals
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.contains
+import kotlin.collections.filter
+
+@OptIn(LLFirInternals::class, SymbolInternals::class)
+class KotlinStandaloneDirectInheritorsProvider(private val project: Project) : KotlinDirectInheritorsProvider {
+    private val staticDeclarationProviderFactory by lazy {
+        KotlinDeclarationProviderFactory.getInstance(project) as? KotlinStaticDeclarationProviderFactory
+            ?: error(
+                "`${KotlinStandaloneDirectInheritorsProvider::class.simpleName}` expects the following declaration provider factory to be" +
+                        " registered: `${KotlinStaticDeclarationProviderFactory::class.simpleName}`"
+            )
+    }
+
+    override fun getDirectKotlinInheritors(
+        ktClass: KtClass,
+        scope: GlobalSearchScope,
+        includeLocalInheritors: Boolean,
+    ): Iterable<KtClassOrObject> {
+        val classId = ktClass.getClassId() ?: return emptyList()
+        val possibleInheritors = staticDeclarationProviderFactory.getDirectInheritorCandidates(classId.shortClassName)
+
+        if (possibleInheritors.isEmpty()) {
+            return emptyList()
+        }
+
+        // The index provides candidates from an original module, not dangling files. If we resolve the supertypes of a candidate in the
+        // context of its session, we will resolve to FIR classes from non-dangling, original modules. If `ktClass` is inside a dangling
+        // file, the FIR class for `ktClass` will come from the dangling module. So we'd compare the original FIR class for the supertype
+        // with the dangling FIR class for `ktClass`, resulting in a mismatch. To avoid such incompatible comparisons, we need to resolve
+        // `ktClass` to the original FIR class.
+        //
+        // Note that this means we don't support providing inheritors based on the dangling file yet, for example if an inheritor was added
+        // or removed only in the dangling file.
+        val baseKtModule = when (val ktModule = ProjectStructureProvider.getModule(project, ktClass, contextualModule = null)) {
+            is KtDanglingFileModule -> ktModule.contextModule
+            else -> ktModule
+        }
+
+        val baseFirClass = ktClass.toFirSymbol(classId, baseKtModule)?.fir as? FirClass ?: return emptyList()
+        return possibleInheritors.filter { isValidInheritor(it, baseFirClass, scope, includeLocalInheritors) }
+    }
+
+    private fun isValidInheritor(
+        candidate: KtClassOrObject,
+        baseFirClass: FirClass,
+        scope: GlobalSearchScope,
+        includeLocalInheritors: Boolean,
+    ): Boolean {
+        if (!includeLocalInheritors && candidate.isLocal) {
+            return false
+        }
+
+        if (!scope.contains(candidate)) {
+            return false
+        }
+
+        val candidateClassId = candidate.getClassId() ?: return false
+        val candidateKtModule = ProjectStructureProvider.getModule(project, candidate, contextualModule = null)
+        val candidateFirSymbol = candidate.toFirSymbol(candidateClassId, candidateKtModule) ?: return false
+        val candidateFirClass = candidateFirSymbol.fir as? FirClass ?: return false
+
+        // `KotlinDirectInheritorsProvider`'s interface guarantees that `getDirectKotlinInheritors` is only called from lazy resolution to
+        // `SEALED_CLASS_INHERITORS` or later, so `isSubClassOf` resolving to `SUPER_TYPES` is legal.
+        return isSubClassOf(candidateFirClass, baseFirClass, allowIndirectSubtyping = false)
+    }
+
+    private fun KtClassOrObject.toFirSymbol(classId: ClassId, ktModule: KtModule): FirClassLikeSymbol<*>? {
+        // Using a resolve session/source-preferred session will cause class stubs from binary libraries to be AST-loaded in IDE mode tests,
+        // which results in an exception since we don't have a decompiler for them. See KT-64898, KT-64899, and KT-64900. If not for these
+        // issues, we would be able to use `analyze` instead of custom session logic.
+        val session = LLFirSessionCache.getInstance(project).getSession(ktModule, preferBinary = true)
+        return session.symbolProvider.getClassLikeSymbolByClassId(classId)
+    }
+}

analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt

@@ -10,7 +10,6 @@ import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.Application
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
-import com.intellij.openapi.vfs.impl.jar.CoreJarFileSystem
 import com.intellij.psi.PsiFile
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analysis.api.standalone.base.project.structure.FirStandaloneServiceRegistrar