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/entrypoint/src/org/jetbrains/kotlin/fir/session/FirAbstractSessionFactory.kt

@@ -13,10 +13,13 @@ import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
 import org.jetbrains.kotlin.fir.extensions.FirSwitchableExtensionDeclarationsSymbolProvider
 import org.jetbrains.kotlin.fir.java.FirCliSession
 import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
-import org.jetbrains.kotlin.fir.resolve.providers.*
-import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCompositeSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY
+import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirCachingCompositeSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirLibrarySessionProvider
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirProviderImpl
+import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.incremental.components.EnumWhenTracker
 import org.jetbrains.kotlin.incremental.components.LookupTracker
@@ -55,7 +58,7 @@ abstract class FirAbstractSessionFactory {
 
             val providers = createProviders(this, builtinsModuleData, kotlinScopeProvider)
 
-            val symbolProvider = FirCompositeSymbolProvider(this, providers)
+            val symbolProvider = FirCachingCompositeSymbolProvider(this, providers)
             register(FirSymbolProvider::class, symbolProvider)
             register(FirProvider::class, FirLibrarySessionProvider(symbolProvider))
         }
@@ -111,10 +114,16 @@ abstract class FirAbstractSessionFactory {
                 dependencyProviders,
             )
 
-            register(FirSymbolProvider::class, FirCompositeSymbolProvider(this, providers))
+            register(
+                FirSymbolProvider::class,
+                FirCachingCompositeSymbolProvider(
+                    this, providers,
+                    expectedCachesToBeCleanedOnce = generatedSymbolsProvider != null
+                )
+            )
 
             generatedSymbolsProvider?.let { register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it) }
-            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, FirCompositeSymbolProvider(this, dependencyProviders))
+            register(DEPENDENCIES_SYMBOL_PROVIDER_QUALIFIED_KEY, FirCachingCompositeSymbolProvider(this, dependencyProviders))
         }
     }
 
@@ -139,7 +148,7 @@ abstract class FirAbstractSessionFactory {
         fun FirSymbolProvider.collectProviders() {
             if (!visited.add(this)) return
             when {
-                this is FirCompositeSymbolProvider -> {
+                this is FirCachingCompositeSymbolProvider -> {
                     for (provider in providers) {
                         provider.collectProviders()
                     }

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

@@ -6,6 +6,9 @@
 package org.jetbrains.kotlin.fir.session
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.deserialization.*
 import org.jetbrains.kotlin.fir.isNewPlaceForBodyGeneration
@@ -14,13 +17,16 @@ import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.library.metadata.KlibDeserializedContainerSource
 import org.jetbrains.kotlin.library.metadata.KlibMetadataClassDataFinder
+import org.jetbrains.kotlin.library.metadata.KlibMetadataProtoBuf
 import org.jetbrains.kotlin.library.metadata.KlibMetadataSerializerProtocol
 import org.jetbrains.kotlin.library.metadata.resolver.KotlinResolvedLibrary
 import org.jetbrains.kotlin.metadata.ProtoBuf
+import org.jetbrains.kotlin.metadata.deserialization.NameResolver
 import org.jetbrains.kotlin.metadata.deserialization.NameResolverImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
+import org.jetbrains.kotlin.serialization.deserialization.getClassId
 import org.jetbrains.kotlin.utils.SmartList
 import java.nio.file.Paths
 
@@ -96,18 +102,63 @@ class KlibBasedSymbolProvider(
 
     override fun computePackageSetWithNonClassDeclarations(): Set<String> = fragmentNamesInLibraries.keys
 
-    // Looks like it's expensive to compute the presence of a class properly for KLib
-    override fun mayHaveTopLevelClass(classId: ClassId): Boolean = true
+    override fun knownTopLevelClassesInPackage(packageFqName: FqName): Set<String> =
+        buildSet {
+            forEachFragmentInPackage(packageFqName) { _, fragment, nameResolver ->
+                for (classNameId in fragment.getExtension(KlibMetadataProtoBuf.className).orEmpty()) {
+                    add(nameResolver.getClassId(classNameId).shortClassName.asString())
+                }
+            }
+        }
 
     @OptIn(SymbolInternals::class)
     override fun extractClassMetadata(classId: ClassId, parentContext: FirDeserializationContext?): ClassMetadataFindResult? {
-        val packageStringName = classId.packageFqName.asString()
+        forEachFragmentInPackage(classId.packageFqName) { resolvedLibrary, fragment, nameResolver ->
+            val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
+            val classProto = finder.findClassData(classId)?.classProto ?: return@forEachFragmentInPackage
 
-        val librariesWithFragment = fragmentNamesInLibraries[packageStringName] ?: return null
+            val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
+            val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return null
+
+            return ClassMetadataFindResult.NoMetadata { symbol ->
+                val source = createDeserializedContainerSource(resolvedLibrary,
+
+                    classId.packageFqName
+                )
+
+                deserializeClassToSymbol(
+                    classId,
+                    classProto,
+                    symbol,
+                    nameResolver,
+                    session,
+                    moduleData,
+                    annotationDeserializer,
+                    kotlinScopeProvider,
+                    KlibMetadataSerializerProtocol,
+                    parentContext,
+                    source,
+                    origin = defaultDeserializationOrigin,
+                    deserializeNestedClass = this::getClass,
+                )
+                symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
+            }
+        }
+
+        return null
+    }
+
+    private inline fun forEachFragmentInPackage(
+        packageFqName: FqName,
+        f: (KotlinResolvedLibrary, ProtoBuf.PackageFragment, NameResolver) -> Unit
+    ) {
+        val packageStringName = packageFqName.asString()
+
+        val librariesWithFragment = fragmentNamesInLibraries[packageStringName] ?: return
 
         for (resolvedLibrary in librariesWithFragment) {
             for (packageMetadataPart in resolvedLibrary.library.packageMetadataParts(packageStringName)) {
-                val libraryPath = Paths.get(resolvedLibrary.library.libraryFile.path)
+
                 val fragment = getPackageFragment(resolvedLibrary, packageStringName, packageMetadataPart)
 
                 val nameResolver = NameResolverImpl(
@@ -115,35 +166,9 @@ class KlibBasedSymbolProvider(
                     fragment.qualifiedNames,
                 )
 
-                val finder = KlibMetadataClassDataFinder(fragment, nameResolver)
-                val classProto = finder.findClassData(classId)?.classProto ?: continue
-
-                val moduleData = moduleDataProvider.getModuleData(libraryPath) ?: return null
-
-                return ClassMetadataFindResult.NoMetadata { symbol ->
-                    val source = createDeserializedContainerSource(resolvedLibrary, classId.packageFqName)
-
-                    deserializeClassToSymbol(
-                        classId,
-                        classProto,
-                        symbol,
-                        nameResolver,
-                        session,
-                        moduleData,
-                        annotationDeserializer,
-                        kotlinScopeProvider,
-                        KlibMetadataSerializerProtocol,
-                        parentContext,
-                        source,
-                        origin = defaultDeserializationOrigin,
-                        deserializeNestedClass = this::getClass,
-                    )
-                    symbol.fir.isNewPlaceForBodyGeneration = isNewPlaceForBodyGeneration(classProto)
-                }
+                f(resolvedLibrary, fragment, nameResolver)
             }
         }
-
-        return null
     }
 
     private fun createDeserializedContainerSource(