[Analysis API] do not preresolve all FirDeclarations to TYPES in tests

Previously, we resolved all declaration to the TYPES phase to collect their sealed inheritors
and thus not really checking lazy resolution to TYPES.
Now sealed class inheritors collection happens in separate LLFirResolveSession

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

@@ -47,6 +47,7 @@ abstract class AbstractCompilerBasedTestForFir : AbstractCompilerBasedTest() {
         configureTest()
         defaultConfiguration(this)
         registerAnalysisApiBaseTestServices(disposable, FirLowLevelCompilerBasedTestConfigurator)
+        usePreAnalysisHandlers(::SealedClassesInheritorsCaclulatorPreAnalysisHandler)
 
         firHandlersStep {
             useHandlers(::LLDiagnosticParameterChecker)

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

@@ -30,7 +30,6 @@ class LowLevelFirAnalyzerFacade(
         get() = ScopeSession()
 
     override fun runCheckers(): Map<FirFile, List<KtDiagnostic>> {
-        findSealedInheritors()
         return allFirFiles.values.associateWith { firFile ->
             val ktFile = firFile.psi as KtFile
             val diagnostics = ktFile.collectDiagnosticsForFile(firResolveSession, diagnosticCheckerFilter)
@@ -39,14 +38,6 @@ class LowLevelFirAnalyzerFacade(
         }
     }
 
-    private fun findSealedInheritors() {
-        allFirFiles.values.forEach { firFile ->
-            firFile.ensureResolved(FirResolvePhase.SUPER_TYPES)
-        }
-        val sealedProcessor = FirSealedClassInheritorsProcessor(allFirFiles.values.first().moduleData.session, ScopeSession())
-        sealedProcessor.process(allFirFiles.values)
-    }
-
     override fun runResolution(): List<FirFile> = shouldNotBeCalled()
     override fun convertToIr(fir2IrExtensions: Fir2IrExtensions): Fir2IrResult = shouldNotBeCalled()
 }

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

@@ -5,15 +5,19 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.compiler.based
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.FirSealedClassInheritorsProcessorFactory
+import org.jetbrains.kotlin.analysis.low.level.api.fir.createFirResolveSessionForNoCaching
+import org.jetbrains.kotlin.analysis.low.level.api.fir.services.LLFirSealedClassInheritorsProcessorFactoryForTests
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.ktModuleProvider
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.FirSealedClassInheritorsProcessor
 import org.jetbrains.kotlin.fir.symbols.ensureResolved
+import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.services.PreAnalysisHandler
 import org.jetbrains.kotlin.test.services.TestModuleStructure
@@ -27,30 +31,37 @@ class SealedClassesInheritorsCaclulatorPreAnalysisHandler(
     override fun preprocessModuleStructure(moduleStructure: TestModuleStructure) {
     }
 
+    // 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
+    // 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) {
         val ktFilesByModule = moduleStructure.modules.associateWith { testModule ->
             testServices.ktModuleProvider.getModuleFiles(testModule).filterIsInstance<KtFile>()
         }
 
-        ktFilesByModule.forEach { (testModule, ktFiles) ->
+        for ((testModule, ktFiles) in ktFilesByModule) {
+            if (ktFiles.isEmpty()) continue
             val project = testServices.compilerConfigurationProvider.getProject(testModule)
-            // 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
-            // preprocessing of all files. Therefore, the IDE does not rely on such a pre-analysis pass of all files in the module.
             val ktModuleProvider = project.getService(ProjectStructureProvider::class.java)
-            val allFirFiles = mutableListOf<FirFile>()
-            ktFiles.forEach { ktFile ->
-                val ktModule = ktModuleProvider.getKtModuleForKtElement(ktFile)
-                val firResolveSession = ktModule.getFirResolveSession(project)
-                allFirFiles.add(ktFile.getOrBuildFirFile(firResolveSession))
-            }
-            allFirFiles.groupBy { it.moduleData.session }.forEach { (session, firFiles) ->
-                for (firFile in firFiles) {
-                    firFile.ensureResolved(FirResolvePhase.TYPES)
-                }
-                val sealedProcessor = FirSealedClassInheritorsProcessor(session, ScopeSession())
-                sealedProcessor.process(firFiles)
-            }
+            val ktModule = ktFiles.map(ktModuleProvider::getKtModuleForKtElement).distinct().single()
+
+            val tmpFirResolveSession = createFirResolveSessionForNoCaching(ktModule, project)
+            val firFiles = ktFiles.map { it.getOrBuildFirFile(tmpFirResolveSession) }
+            val sealedInheritors = collectSealedClassInheritors(firFiles, tmpFirResolveSession)
+            val provider =
+                project.getService(FirSealedClassInheritorsProcessorFactory::class.java) as LLFirSealedClassInheritorsProcessorFactoryForTests
+            provider.registerInheritors(ktModule, sealedInheritors)
         }
     }
+
+    private fun collectSealedClassInheritors(
+        firFiles: List<FirFile>,
+        tmpFirResolveSession: LLFirResolveSession
+    ): Map<ClassId, List<ClassId>> {
+        firFiles.forEach { it.ensureResolved(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 }
+    }
 }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/services/LLFirSealedClassInheritorsProcessorFactoryForTests.kt

@@ -6,11 +6,32 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.services
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.FirSealedClassInheritorsProcessorFactory
+import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProvider
-import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProviderImpl
+import org.jetbrains.kotlin.fir.declarations.utils.classId
+import org.jetbrains.kotlin.name.ClassId
 
 internal class LLFirSealedClassInheritorsProcessorFactoryForTests : FirSealedClassInheritorsProcessorFactory() {
-    override fun createSealedClassInheritorsProvider(): SealedClassInheritorsProvider {
-        return SealedClassInheritorsProviderImpl
+    private val inheritorsByModule = mutableMapOf<KtModule, Map<ClassId, List<ClassId>>>()
+
+    fun registerInheritors(ktModule: KtModule, inheritors: Map<ClassId, List<ClassId>>) {
+        inheritorsByModule[ktModule] = inheritors
     }
+
+    override fun createSealedClassInheritorsProvider(): SealedClassInheritorsProvider {
+        return SealedClassInheritorsProviderForTests(inheritorsByModule)
+    }
+}
+
+private class SealedClassInheritorsProviderForTests(
+    private val inheritorsByModule: Map<KtModule, Map<ClassId, List<ClassId>>>
+) : SealedClassInheritorsProvider() {
+    override fun getSealedClassInheritors(firClass: FirRegularClass): List<ClassId> {
+        val ktModule = firClass.llFirModuleData.ktModule
+        val inheritorsForModuleMap = inheritorsByModule.getValue(ktModule)
+        return inheritorsForModuleMap[firClass.classId].orEmpty()
+    }
+
 }

compiler/test-infrastructure/tests/org/jetbrains/kotlin/test/TestRunner.kt

@@ -81,6 +81,10 @@ class TestRunner(private val testConfiguration: TestConfiguration) {
             preprocessor.preprocessModuleStructure(moduleStructure)
         }
 
+        testConfiguration.preAnalysisHandlers.forEach { preprocessor ->
+            preprocessor.prepareSealedClassInheritors(moduleStructure)
+        }
+
         for (module in modules) {
             val shouldProcessNextModules = processModule(module, dependencyProvider)
             if (!shouldProcessNextModules) break

compiler/testData/diagnostics/tests/multimodule/kt14249.kt

@@ -1,4 +1,3 @@
-// FIR_IDE_IGNORE
 // FIR_IDENTICAL
 // MODULE: m1
 // FILE: test/Foo.java