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

2021-04-01 17:55:24 +02:00
parent 0e93924ff3
commit 5ceec858aa
13 changed files with 133 additions and 117 deletions
@@ -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 --------------------------

@@ -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)

@@ -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)
@@ -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))
            }
@@ -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) {
@@ -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()
@@ -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)
            }
@@ -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
@@ -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)
-}
@@ -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)
@@ -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)
+    }
+}
@@ -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())
    }
 }
@@ -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"))
    }
 }