FIR IDE: do not search for sealed class inheritors on every class resolve

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/ComponentsContainers.kt

@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.fir.analysis.CheckersComponent
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirNameConflictsTracker
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import org.jetbrains.kotlin.fir.caches.FirThreadUnsafeCachesFactory
+import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProvider
+import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProviderImpl
 import org.jetbrains.kotlin.fir.extensions.FirExtensionService
 import org.jetbrains.kotlin.fir.extensions.FirPredicateBasedProvider
 import org.jetbrains.kotlin.fir.extensions.FirRegisteredPluginAnnotations
@@ -51,6 +53,12 @@ fun FirSession.registerThreadUnsafeCaches() {
     register(FirCachesFactory::class, FirThreadUnsafeCachesFactory)
 }
 
+@OptIn(SessionConfiguration::class)
+fun FirSession.registerSealedClassInheritorsProvider() {
+    register(SealedClassInheritorsProvider::class, SealedClassInheritorsProviderImpl)
+}
+
+
 
 // -------------------------- Resolve components --------------------------
 

compiler/fir/entrypoint/src/org/jetbrains/kotlin/fir/session/FirSessionFactory.kt

@@ -77,6 +77,7 @@ object FirSessionFactory {
             registerCommonComponents(languageVersionSettings)
             registerResolveComponents(lookupTracker)
             registerJavaSpecificResolveComponents()
+            registerSealedClassInheritorsProvider()
 
             val kotlinScopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
 
@@ -127,6 +128,7 @@ object FirSessionFactory {
         return FirLibrarySession(moduleInfo, sessionProvider).apply {
             registerThreadUnsafeCaches()
             registerCommonComponents(languageVersionSettings)
+            registerSealedClassInheritorsProvider()
 
             val javaSymbolProvider = JavaSymbolProvider(this, project, scope)
 

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/deserialization/ClassDeserialization.kt

@@ -195,9 +195,10 @@ fun deserializeClassToSymbol(
         })
     }.also {
         if (isSealed) {
-            it.sealedInheritors = classProto.sealedSubclassFqNameList.map { nameIndex ->
+            val inheritors = classProto.sealedSubclassFqNameList.map { nameIndex ->
                 ClassId.fromString(nameResolver.getQualifiedClassName(nameIndex))
             }
+            it.setSealedClassInheritors(inheritors)
         }
         (it.annotations as MutableList<FirAnnotationCall>) +=
             context.annotationDeserializer.loadClassAnnotations(classProto, context.nameResolver)

compiler/fir/fir-serialization/src/org/jetbrains/kotlin/fir/serialization/FirElementSerializer.kt

@@ -166,8 +166,7 @@ class FirElementSerializer private constructor(
         }
 
         if (klass is FirRegularClass && klass.modality == Modality.SEALED) {
-            val inheritors = klass.sealedInheritors
-            requireNotNull(inheritors)
+            val inheritors = klass.getSealedClassInheritors(session)
             for (inheritorId in inheritors) {
                 builder.addSealedSubclassFqName(stringTable.getQualifiedClassNameIndex(inheritorId))
             }

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaSymbolProvider.kt

@@ -332,10 +332,11 @@ class JavaSymbolProvider(
             }
         }.apply {
             if (modality == Modality.SEALED) {
-                sealedInheritors = javaClass.permittedTypes.mapNotNull { classifierType ->
+                val inheritors = javaClass.permittedTypes.mapNotNull { classifierType ->
                     val classifier = classifierType.classifier as? JavaClass
                     classifier?.let { JavaToKotlinClassMap.mapJavaToKotlin(it.fqName!!) }
                 }
+                setSealedClassInheritors(inheritors)
             }
 
             if (classIsAnnotation) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirSealedClassInheritorsProcessor.kt

@@ -86,7 +86,7 @@ class FirSealedClassInheritorsProcessor(
             if (regularClass.modality == Modality.SEALED) {
                 val inheritors = inheritorsMap.remove(regularClass)
                 if (inheritors != null) {
-                    regularClass.sealedInheritors = inheritors
+                    regularClass.setSealedClassInheritors(inheritors)
                 }
             }
             if (inheritorsMap.isEmpty()) return regularClass.compose()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt

@@ -313,7 +313,7 @@ private object WhenOnSealedClassExhaustivenessChecker : WhenExhaustivenessChecke
             return
         }
         when {
-            fir.modality == Modality.SEALED -> fir.sealedInheritors?.forEach {
+            fir.modality == Modality.SEALED -> fir.getSealedClassInheritors(session).forEach {
                 val symbol = session.symbolProvider.getClassLikeSymbolByFqName(it) as? FirRegularClassSymbol
                 symbol?.collectAllSubclassesTo(destination, session)
             }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/declarations/SealedClassInheritors.kt

@@ -5,8 +5,41 @@
 
 package org.jetbrains.kotlin.fir.declarations
 
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.name.ClassId
 
-object SealedClassInheritorsKey : FirDeclarationDataKey()
 
-var FirRegularClass.sealedInheritors: List<ClassId>? by FirDeclarationDataRegistry.data(SealedClassInheritorsKey)
+@RequiresOptIn("For getting/setting sealed class inheritors, consider using getSealedClassInheritors/setSealedClassInheritors")
+annotation class SealedClassInheritorsProviderInternals
+
+abstract class SealedClassInheritorsProvider : FirSessionComponent {
+    abstract fun getSealedClassInheritors(firClass: FirRegularClass): List<ClassId>
+}
+
+private val FirSession.sealedClassInheritorsProvider: SealedClassInheritorsProvider by FirSession.sessionComponentAccessor()
+
+object SealedClassInheritorsProviderImpl : SealedClassInheritorsProvider() {
+    @OptIn(SealedClassInheritorsProviderInternals::class)
+    override fun getSealedClassInheritors(firClass: FirRegularClass): List<ClassId> {
+        return firClass.sealedInheritorsAttr ?: emptyList()
+    }
+}
+
+
+fun FirRegularClass.getSealedClassInheritors(session: FirSession): List<ClassId> {
+    require(this.isSealed)
+    return session.sealedClassInheritorsProvider.getSealedClassInheritors(this)
+}
+
+@OptIn(SealedClassInheritorsProviderInternals::class)
+fun FirRegularClass.setSealedClassInheritors(inheritors: List<ClassId>) {
+    require(this.isSealed)
+    sealedInheritorsAttr = inheritors
+}
+
+private object SealedClassInheritorsKey : FirDeclarationDataKey()
+
+@SealedClassInheritorsProviderInternals
+var FirRegularClass.sealedInheritorsAttr: List<ClassId>? by FirDeclarationDataRegistry.data(SealedClassInheritorsKey)
+    private set