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
+            val inheritors = klass.getSealedClassInheritors(session)
-            requireNotNull(inheritors)
             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

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirIdeSealedHierarchyProcessor.kt

@@ -1,104 +0,0 @@
-/*
- * 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.fir.low.level.api
-
-import com.intellij.openapi.module.Module
-import com.intellij.psi.JavaDirectoryService
-import com.intellij.psi.PsiPackage
-import com.intellij.psi.search.GlobalSearchScope
-import com.intellij.psi.search.PackageScope
-import com.intellij.psi.search.SearchScope
-import com.intellij.psi.search.searches.ClassInheritorsSearch
-import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
-import org.jetbrains.kotlin.asJava.toLightClass
-import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
-import org.jetbrains.kotlin.fir.resolve.transformers.FirSealedClassInheritorsProcessor
-import org.jetbrains.kotlin.fir.resolve.transformers.FirTransformerBasedResolveProcessor
-import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
-import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
-import org.jetbrains.kotlin.fir.visitors.FirTransformer
-import org.jetbrains.kotlin.fir.visitors.compose
-import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
-import org.jetbrains.kotlin.idea.util.module
-import org.jetbrains.kotlin.name.ClassId
-import org.jetbrains.kotlin.psi.KtClass
-import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
-
-
-class FirIdeSealedHierarchyProcessor(session: FirSession, scopeSession: ScopeSession) :
-    FirTransformerBasedResolveProcessor(session, scopeSession) {
-
-    override val transformer: FirTransformer<Nothing?> = SealedClassInheritorsTransformer
-
-
-    private object SealedClassInheritorsTransformer : FirTransformer<Nothing?>() {
-        override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
-            throw IllegalStateException("Should not be there")
-        }
-
-        override fun transformFile(file: FirFile, data: Nothing?): CompositeTransformResult<FirDeclaration> {
-            val sealedClassInheritorsMap = mutableMapOf<FirRegularClass, MutableList<ClassId>>()
-            file.accept(SealedInheritorsCollector, sealedClassInheritorsMap)
-
-            if (sealedClassInheritorsMap.isEmpty()) return file.compose()
-            return file.transform(FirSealedClassInheritorsProcessor.InheritorsTransformer(sealedClassInheritorsMap), null)
-        }
-    }
-
-    private object SealedInheritorsCollector : FirDefaultVisitor<Unit, MutableMap<FirRegularClass, MutableList<ClassId>>>() {
-
-        override fun visitElement(element: FirElement, data: MutableMap<FirRegularClass, MutableList<ClassId>>) {}
-
-        override fun visitFile(file: FirFile, data: MutableMap<FirRegularClass, MutableList<ClassId>>) {
-            file.declarations.forEach { it.accept(this, data) }
-        }
-
-        override fun visitRegularClass(regularClass: FirRegularClass, data: MutableMap<FirRegularClass, MutableList<ClassId>>) {
-            if (!regularClass.isSealed) {
-                regularClass.acceptChildren(this, data)
-                return
-            }
-
-            val sealedKtClass = regularClass.psi as? KtClass ?: return
-            val module = sealedKtClass.module ?: return
-            val containingPackage = regularClass.classId.packageFqName
-
-            val psiPackage = KotlinJavaPsiFacade.getInstance(sealedKtClass.project)
-                .findPackage(containingPackage.asString(), GlobalSearchScope.moduleScope(module))
-                ?: getPackageViaDirectoryService(sealedKtClass)
-                ?: return
-
-            val kotlinAsJavaSupport = KotlinAsJavaSupport.getInstance(sealedKtClass.project)
-            val lightClass = sealedKtClass.toLightClass() ?: kotlinAsJavaSupport.getFakeLightClass(sealedKtClass)
-
-            val searchScope: SearchScope = getSearchScope(module, psiPackage)
-            val searchParameters = ClassInheritorsSearch.SearchParameters(lightClass, searchScope, false, true, false)
-            val subclasses = ClassInheritorsSearch.search(searchParameters)
-                .mapNotNull { it.classIdIfNonLocal() }
-                .toMutableList()
-
-            // Enforce a deterministic order on the result.
-            subclasses.sortBy { it.toString() }
-
-            data[regularClass] = subclasses
-        }
-
-        private fun getSearchScope(module: Module, psiPackage: PsiPackage): GlobalSearchScope {
-            val packageScope = PackageScope(psiPackage, false, false)
-            // MPP multiple common modules are not supported!!
-            return module.moduleScope.intersectWith(packageScope)
-        }
-    }
-}
-
-private fun getPackageViaDirectoryService(ktClass: KtClass): PsiPackage? {
-    val directory = ktClass.containingFile.containingDirectory ?: return null
-    return JavaDirectoryService.getInstance().getPackage(directory)
-}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirPhaseRunner.kt

@@ -81,7 +81,7 @@ internal fun FirResolvePhase.createTransformerBasedProcessorByPhase(
         FirResolvePhase.CLASS_GENERATION -> FirDummyTransformerBasedProcessor(session, scopeSession) // TODO: remove
         FirResolvePhase.IMPORTS -> FirImportResolveProcessor(session, scopeSession)
         FirResolvePhase.SUPER_TYPES -> FirSupertypeResolverProcessor(session, scopeSession)
-        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirIdeSealedHierarchyProcessor(session, scopeSession)
+        FirResolvePhase.SEALED_CLASS_INHERITORS -> FirDummyTransformerBasedProcessor(session, scopeSession)
         FirResolvePhase.TYPES -> FirTypeResolveProcessor(session, scopeSession)
         FirResolvePhase.ARGUMENTS_OF_PLUGIN_ANNOTATIONS -> FirAnnotationArgumentsResolveProcessor(session, scopeSession)
         FirResolvePhase.EXTENSION_STATUS_UPDATE -> FirTransformerBasedExtensionStatusProcessor(session, scopeSession)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/SealedClassInheritorsProviderIdeImpl.kt

@@ -0,0 +1,70 @@
+/*
+ * 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.fir.low.level.api
+
+import com.intellij.openapi.module.Module
+import com.intellij.psi.JavaDirectoryService
+import com.intellij.psi.PsiPackage
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.PackageScope
+import com.intellij.psi.search.SearchScope
+import com.intellij.psi.search.searches.ClassInheritorsSearch
+import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
+import org.jetbrains.kotlin.asJava.toLightClass
+import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.psi
+import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
+import org.jetbrains.kotlin.idea.util.module
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.resolve.jvm.KotlinJavaPsiFacade
+import java.util.concurrent.ConcurrentHashMap
+
+class SealedClassInheritorsProviderIdeImpl : SealedClassInheritorsProvider() {
+    val cache = ConcurrentHashMap<ClassId, List<ClassId>>()
+
+    @OptIn(SealedClassInheritorsProviderInternals::class)
+    override fun getSealedClassInheritors(firClass: FirRegularClass): List<ClassId> {
+        require(firClass.isSealed)
+        firClass.sealedInheritorsAttr?.let { return it }
+        return cache.computeIfAbsent(firClass.classId) { getInheritors(firClass) }
+    }
+
+    private fun getInheritors(firClass: FirRegularClass): List<ClassId> {
+        val sealedKtClass = firClass.psi as? KtClass ?: return emptyList()
+        val module = sealedKtClass.module ?: return emptyList()
+        val containingPackage = firClass.classId.packageFqName
+
+        val psiPackage = KotlinJavaPsiFacade.getInstance(sealedKtClass.project)
+            .findPackage(containingPackage.asString(), GlobalSearchScope.moduleScope(module))
+            ?: getPackageViaDirectoryService(sealedKtClass)
+            ?: return emptyList()
+
+        val kotlinAsJavaSupport = KotlinAsJavaSupport.getInstance(sealedKtClass.project)
+        val lightClass = sealedKtClass.toLightClass() ?: kotlinAsJavaSupport.getFakeLightClass(sealedKtClass)
+
+        val searchScope: SearchScope = getSearchScope(module, psiPackage)
+        val searchParameters = ClassInheritorsSearch.SearchParameters(lightClass, searchScope, false, true, false)
+        val subclasses = ClassInheritorsSearch.search(searchParameters)
+            .mapNotNull { it.classIdIfNonLocal() }
+            .toMutableList()
+
+        // Enforce a deterministic order on the result.
+        subclasses.sortBy { it.toString() }
+        return subclasses
+    }
+
+    private fun getSearchScope(module: Module, psiPackage: PsiPackage): GlobalSearchScope {
+        val packageScope = PackageScope(psiPackage, false, false)
+        // MPP multiple common modules are not supported!!
+        return module.moduleScope.intersectWith(packageScope)
+    }
+
+    private fun getPackageViaDirectoryService(ktClass: KtClass): PsiPackage? {
+        val directory = ktClass.containingFile.containingDirectory ?: return null
+        return JavaDirectoryService.getInstance().getPackage(directory)
+    }
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSessionFactory.kt

@@ -13,15 +13,14 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor
 import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.backend.jvm.FirJvmTypeMapper
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
-import org.jetbrains.kotlin.fir.checkers.registerCommonCheckers
 import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
+import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProvider
 import org.jetbrains.kotlin.fir.dependenciesWithoutSelf
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinDeserializedJvmSymbolsProvider
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.providers.impl.FirBuiltinSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCloneableSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
@@ -33,6 +32,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.IDEPackagePartProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.IdeFirPhaseManager
 import org.jetbrains.kotlin.idea.fir.low.level.api.IdeSessionComponents
+import org.jetbrains.kotlin.idea.fir.low.level.api.SealedClassInheritorsProviderIdeImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCacheImpl
 import org.jetbrains.kotlin.idea.fir.low.level.api.fir.caches.FirThreadSafeCachesFactory
@@ -43,7 +43,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeLibrariesSess
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirModuleWithDependenciesSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.FirThreadSafeSymbolProviderWrapper
-import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionFactory.registerIdeComponents
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.ModuleLibrariesSearchScope
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.load.java.JavaClassFinderImpl
@@ -232,5 +231,6 @@ internal object FirIdeSessionFactory {
     private fun FirIdeSession.registerIdeComponents() {
         register(IdeSessionComponents::class, IdeSessionComponents.create(this))
         register(FirCachesFactory::class, FirThreadSafeCachesFactory)
+        register(SealedClassInheritorsProvider::class, SealedClassInheritorsProviderIdeImpl())
     }
 }

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/AbstractFirSealedInheritorsTest.kt

@@ -25,7 +25,13 @@ abstract class AbstractFirSealedInheritorsTest : AbstractFirMultiModuleLazyResol
 
     override fun checkFirFile(firFile: FirFile, path: String) {
         val allClasses = firFile.listNestedClasses().closure { it.listNestedClasses() }
-        val inheritorNames = allClasses.flatMap { it.sealedInheritors ?: emptyList() }.map { it.asString() }.sorted()
+        val inheritorNames = allClasses.flatMap { firClass ->
+            if (firClass.isSealed) {
+                firClass.getSealedClassInheritors(firFile.session)
+            } else {
+                emptyList()
+            }
+        }.map { it.asString() }.sorted()
         KotlinTestUtils.assertEqualsToFile(File("$path/expected.txt"), inheritorNames.joinToString("\n"))
     }
 }