[AA] Tests: Rewrite sealed class inheritors collection to use the declaration provider

- Now that binary libraries are decompiled to stubs instead of PSI
  files, we cannot collect sealed inheritors from `KtFile`s anymore.
  Since all `KtFile`s and binary library stubs are both indexed by the
  declaration provider, we can collect inheritors from its index
  instead.
- Invalidating all sessions at the end of `prepareSealedClassInheritors`
  fixes some improper resolve phases in lazy resolution test data. While
  the previous implementation requested an uncached resolve session, it
  didn't account for sessions of dependencies still being cached.

^KT-65960

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

@@ -39,6 +39,7 @@ import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.impl.*
 import org.jetbrains.kotlin.serialization.deserialization.builtins.BuiltInSerializerProtocol
+import org.jetbrains.kotlin.utils.addToStdlib.flattenTo
 import java.util.concurrent.ConcurrentHashMap
 
 public class KotlinStaticDeclarationProvider internal constructor(
@@ -396,6 +397,8 @@ public class KotlinStaticDeclarationProviderFactory(
     public fun getAdditionalCreatedKtFiles(): List<KtFile> {
         return createdFakeKtFiles
     }
+
+    public fun getAllKtClasses(): List<KtClassOrObject> = index.classMap.values.flattenTo(mutableListOf())
 }
 
 /**

analysis/low-level-api-fir/testData/lazyResolve/functions/actualOverride.txt

@@ -1,14 +1,14 @@
 RAW_FIR:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? open expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -44,16 +44,16 @@ FILE: [ResolvedTo(RAW_FIR)] jvm.kt
     }
 
 IMPORTS:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? open expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -89,16 +89,16 @@ FILE: [ResolvedTo(IMPORTS)] jvm.kt
     }
 
 COMPILER_REQUIRED_ANNOTATIONS:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? open expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -134,16 +134,16 @@ FILE: [ResolvedTo(IMPORTS)] jvm.kt
     }
 
 COMPANION_GENERATION:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? open expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -201,8 +201,8 @@ FILE: [ResolvedTo(IMPORTS)] common.kt
     }
 
 FILE: [ResolvedTo(IMPORTS)] jvm.kt
-    public open actual [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
-        public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>| {
+    public open actual [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
+        public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>| {
             LAZY_super<R|kotlin/Any|>
         }
 

analysis/low-level-api-fir/testData/lazyResolve/functions/transitiveSubstitutionFakeOverride.txt

@@ -1,14 +1,14 @@
 RAW_FIR:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? final? expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -46,16 +46,16 @@ FILE: [ResolvedTo(RAW_FIR)] jvm.kt
     }
 
 IMPORTS:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? final? expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -93,16 +93,16 @@ FILE: [ResolvedTo(IMPORTS)] jvm.kt
     }
 
 COMPILER_REQUIRED_ANNOTATIONS:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? final? expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -140,16 +140,16 @@ FILE: [ResolvedTo(IMPORTS)] jvm.kt
     }
 
 COMPANION_GENERATION:
-FILE: [ResolvedTo(IMPORTS)] common.kt
-    public open expect [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
-        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>|
+FILE: [ResolvedTo(RAW_FIR)] common.kt
+    public open expect [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
+        public? expect [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>|
 
         public? final? expect [ResolvedTo(RAW_FIR)] fun existingMethodInBase([ResolvedTo(RAW_FIR)] param: T): R|kotlin/Unit|
 
     }
-    public open [ResolvedTo(SUPER_TYPES)] class Transitive : R|Base<kotlin/String>| {
+    public open [ResolvedTo(RAW_FIR)] class Transitive : Base<String> {
         public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Transitive] constructor(): R|Transitive| {
-            LAZY_super<Base<R|kotlin/String|>>
+            LAZY_super<Base<String>>
         }
 
     }
@@ -209,8 +209,8 @@ FILE: [ResolvedTo(IMPORTS)] common.kt
     }
 
 FILE: [ResolvedTo(IMPORTS)] jvm.kt
-    public open actual [ResolvedTo(RAW_FIR)] class Base<[ResolvedTo(RAW_FIR)] T> : R|kotlin/Any| {
-        public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(RAW_FIR)] T>(): R|Base<T>| {
+    public open actual [ResolvedTo(SUPER_TYPES)] class Base<[ResolvedTo(SUPER_TYPES)] T> : R|kotlin/Any| {
+        public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=Base] constructor<[ResolvedTo(SUPER_TYPES)] T>(): R|Base<T>| {
             LAZY_super<R|kotlin/Any|>
         }
 

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/SealedClassesInheritorsCalculatorPreAnalysisHandler.kt

@@ -5,20 +5,23 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.compiler.based
 
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.LLSealedInheritorsProviderFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLSealedInheritorsProviderFactoryForTests
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
-import org.jetbrains.kotlin.analysis.test.framework.project.structure.getKtFiles
+import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProviderFactory
+import org.jetbrains.kotlin.analysis.providers.impl.KotlinStaticDeclarationProviderFactory
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
-import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.analysis.test.framework.project.structure.mainModules
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSealedClassInheritorsProcessor
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhaseRecursively
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.test.directives.model.DirectiveApplicability
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
 import org.jetbrains.kotlin.test.services.PreAnalysisHandler
@@ -34,40 +37,63 @@ class SealedClassesInheritorsCalculatorPreAnalysisHandler(
     }
 
     // Manually process all inheritors of sealed classes so that SealedClassInheritorsProviderTestImpl can work correctly for tests.
-    // In the actual IDE, SealedClassInheritorsProviderIdeImpl works by finding inheritors from the index instead of do a
+    // In the actual IDE, SealedClassInheritorsProviderIdeImpl works by finding inheritors from the index instead of doing a
     // preprocessing of all files. Therefore, the IDE does not rely on such a pre-analysis pass of all files in the module.
     override fun prepareSealedClassInheritors(moduleStructure: TestModuleStructure) {
         if (Directives.DISABLE_SEALED_INHERITOR_CALCULATOR in moduleStructure.allDirectives) {
             return
         }
 
-        val ktFilesByModule = moduleStructure.modules.associateWith { testModule ->
-            testServices.ktModuleProvider.getKtFiles(testModule)
+        val project = testServices.compilerConfigurationProvider.getProject(moduleStructure.modules.first())
+
+        val sealedInheritorsProviderFactory =
+            project.getService(LLSealedInheritorsProviderFactory::class.java) as? LLSealedInheritorsProviderFactoryForTests
+                ?: return
+
+        val declarationProviderFactory = KotlinDeclarationProviderFactory.getInstance(project) as KotlinStaticDeclarationProviderFactory
+        val projectStructureProvider = project.getService(ProjectStructureProvider::class.java)
+
+        val allKtClasses = declarationProviderFactory.getAllKtClasses()
+        val ktClassesByKtModule = allKtClasses.groupBy { projectStructureProvider.getModule(it, contextualModule = null) }
+
+        for (ktTestModule in testServices.ktModuleProvider.mainModules) {
+            val ktClasses = ktClassesByKtModule[ktTestModule.ktModule] ?: continue
+            val sealedInheritors = collectSealedInheritorClassIds(ktTestModule.ktModule, ktClasses, project)
+
+            sealedInheritorsProviderFactory.registerInheritors(ktTestModule.ktModule, sealedInheritors)
         }
 
-        for ((testModule, ktFiles) in ktFilesByModule) {
-            if (ktFiles.isEmpty()) continue
-            val project = testServices.compilerConfigurationProvider.getProject(testModule)
-            val projectStructureProvider = project.getService(ProjectStructureProvider::class.java)
-            val ktModule = ktFiles.map { projectStructureProvider.getModule(it, contextualModule = null) }.distinct().single()
-
-            val tmpFirResolveSession = LLFirResolveSessionService.getInstance(project).getFirResolveSessionNoCaching(ktModule)
-            val firFiles = ktFiles.map { it.getOrBuildFirFile(tmpFirResolveSession) }
-            val sealedInheritors = collectSealedClassInheritors(firFiles, tmpFirResolveSession)
-            val provider = project.getService(LLSealedInheritorsProviderFactory::class.java) as LLSealedInheritorsProviderFactoryForTests
-            provider.registerInheritors(ktModule, sealedInheritors)
-        }
+        // We request and cache sessions while collecting sealed inheritors. To not interfere with the test, we should invalidate all
+        // sessions. Note that, while we could get uncached sessions from `LLFirSessionCache`, dependency sessions will still be cached, so
+        // that is not really an alternative to invalidating sessions.
+        LLFirSessionCache.getInstance(project).removeAllSessions(includeLibraryModules = true)
     }
 
-    private fun collectSealedClassInheritors(
-        firFiles: List<FirFile>,
-        tmpFirResolveSession: LLFirResolveSession,
+    private fun collectSealedInheritorClassIds(
+        ktModule: KtModule,
+        ktClasses: List<KtClassOrObject>,
+        project: Project,
     ): Map<ClassId, List<ClassId>> {
-        firFiles.forEach { it.lazyResolveToPhaseRecursively(FirResolvePhase.TYPES) }
-        val inheritorsCollector = FirSealedClassInheritorsProcessor.InheritorsCollector(tmpFirResolveSession.useSiteFirSession)
-        val sealedClassInheritorsMap = mutableMapOf<FirRegularClass, MutableList<ClassId>>()
-        firFiles.forEach { it.accept(inheritorsCollector, sealedClassInheritorsMap) }
-        return sealedClassInheritorsMap.mapKeys { (firClass, _) -> firClass.symbol.classId }
+        val sealedInheritorsByFirClass = mutableMapOf<FirRegularClass, MutableList<ClassId>>()
+
+        ktClasses.forEach { ktClass ->
+            val classId = ktClass.getClassId() ?: return@forEach
+
+            // Using a resolve session/source-preferred session will cause class stubs from binary libraries to be decompiled (which
+            // results in an exception since we don't have a decompiler for them).
+            val firSession = LLFirSessionCache.getInstance(project).getSession(ktModule, preferBinary = true)
+            val firClassSymbol = firSession.symbolProvider.getClassLikeSymbolByClassId(classId) ?: return@forEach
+            firClassSymbol.lazyResolveToPhaseRecursively(FirResolvePhase.TYPES)
+
+            val inheritorsCollector = FirSealedClassInheritorsProcessor.InheritorsCollector(firSession)
+            firClassSymbol.fir.accept(inheritorsCollector, sealedInheritorsByFirClass)
+        }
+
+        return buildMap {
+            sealedInheritorsByFirClass.forEach { firClass, classIds ->
+                put(firClass.symbol.classId, classIds.distinct())
+            }
+        }
     }
 
     object Directives : SimpleDirectivesContainer() {