K2: Introduce FirCachingCompositeSymbolProvider

The main idea is pre-computing the sets of names that might be
met there, that helps to decrease the sizes of the backing maps
(by avoiding irrelevant keys)

Totally, this branch with previous commits speeds up MT Full Kotlin
approximately on 3 seconds (~5%)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirCloneableSymbolProvider.kt

@@ -77,4 +77,13 @@ class FirCloneableSymbolProvider(
     override fun getPackage(fqName: FqName): FqName? {
         return null
     }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String> = emptySet()
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
+        if (packageFqName == StandardClassIds.Cloneable.packageFqName)
+            setOf(StandardClassIds.Cloneable.shortClassName.asString())
+        else
+            emptySet()
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> = emptySet()
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirProviderImpl.kt

@@ -75,6 +75,26 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
             if (fqName in state.allSubPackages) return fqName
             return null
         }
+
+        override fun computePackageSetWithTopLevelCallables(): Set<String> =
+            state.allSubPackages.mapTo(mutableSetOf()) { it.asString() }
+
+        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
+            state.classifierInPackage[packageFqName].orEmpty().mapTo(mutableSetOf()) { it.asString() }
+
+        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> = buildSet {
+            for (key in state.functionMap.keys) {
+                if (key.packageName == packageFqName) {
+                    add(key.callableName)
+                }
+            }
+
+            for (key in state.propertyMap.keys) {
+                if (key.packageName == packageFqName) {
+                    add(key.callableName)
+                }
+            }
+        }
     }
 
     private val FirDeclaration.file: FirFile
@@ -109,6 +129,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
 
             if (!classId.isNestedClass && !classId.isLocal) {
                 data.state.classesInPackage.getOrPut(classId.packageFqName, ::mutableSetOf).add(classId.shortClassName)
+                data.state.classifierInPackage.getOrPut(classId.packageFqName, ::mutableSetOf).add(classId.shortClassName)
             }
 
             regularClass.acceptChildren(this, data)
@@ -120,6 +141,8 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
             data.state.classifierMap.put(classId, typeAlias)?.let {
                 data.nameConflictsTracker?.registerClassifierRedeclaration(classId, typeAlias.symbol, data.file, it.symbol, prevFile)
             }
+
+            data.state.classifierInPackage.getOrPut(classId.packageFqName, ::mutableSetOf).add(classId.shortClassName)
         }
 
         override fun visitPropertyAccessor(
@@ -166,10 +189,11 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
     private val state = State()
 
     private class State {
-        val fileMap = mutableMapOf<FqName, List<FirFile>>()
+        val fileMap: MutableMap<FqName, List<FirFile>> = mutableMapOf<FqName, List<FirFile>>()
         val allSubPackages = mutableSetOf<FqName>()
         val classifierMap = mutableMapOf<ClassId, FirClassLikeDeclaration>()
         val classifierContainerFileMap = mutableMapOf<ClassId, FirFile>()
+        val classifierInPackage = mutableMapOf<FqName, MutableSet<Name>>()
         val classesInPackage = mutableMapOf<FqName, MutableSet<Name>>()
         val functionMap = mutableMapOf<CallableId, List<FirNamedFunctionSymbol>>()
         val propertyMap = mutableMapOf<CallableId, List<FirPropertySymbol>>()
@@ -195,6 +219,7 @@ class FirProviderImpl(val session: FirSession, val kotlinScopeProvider: FirKotli
             constructorMap.putAll(other.constructorMap)
             callableContainerMap.putAll(other.callableContainerMap)
             classesInPackage.putAll(other.classesInPackage)
+            classifierInPackage.putAll(other.classifierInPackage)
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirCompilerRequiredAnnotationsResolveTransformer.kt

@@ -7,12 +7,16 @@ package org.jetbrains.kotlin.fir.resolve.transformers.plugin
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.SessionConfiguration
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase.COMPILER_REQUIRED_ANNOTATIONS
 import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCachingCompositeSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.transformers.*
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.LocalClassesNavigationInfo
 import org.jetbrains.kotlin.fir.types.*
@@ -44,7 +48,18 @@ class FirCompilerRequiredAnnotationsResolveProcessor(
     @OptIn(FirSymbolProviderInternals::class)
     override fun afterPhase() {
         super.afterPhase()
-        session.generatedDeclarationsSymbolProvider?.enable()
+
+        val generatedDeclarationsSymbolProvider = session.generatedDeclarationsSymbolProvider
+        generatedDeclarationsSymbolProvider?.enable()
+
+        // This part is a bit hacky way to clear the caches in FirCachingCompositeSymbolProvider when there are plugins that may generate new entities.
+        // It's necessary because otherwise, when symbol provider is being queried on the stage of compiler-required annotations resolution
+        // we record incorrect (incomplete) results to its cache, so after the phase is completed we just start from the scratch
+        val symbolProvider = session.symbolProvider
+        if (generatedDeclarationsSymbolProvider != null && symbolProvider is FirCachingCompositeSymbolProvider) {
+            @OptIn(SessionConfiguration::class)
+            session.register(FirSymbolProvider::class, symbolProvider.createCopyWithCleanCaches())
+        }
     }
 }