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

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirModuleWithDependenciesSymbolProvider.kt

@@ -77,6 +77,12 @@ internal class LLFirModuleWithDependenciesSymbolProvider(
         getPackageWithoutDependencies(fqName)
             ?: dependencyProvider.getPackage(fqName)
 
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
+
 
     fun getPackageWithoutDependencies(fqName: FqName): FqName? =
         providers.firstNotNullOfOrNull { it.getPackage(fqName) }
@@ -150,6 +156,10 @@ internal abstract class LLFirDependentModuleProviders(
             }
         }
 
+    // TODO: Consider having proper implementations for sake of optimizations
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
 
     private fun <S : FirCallableSymbol<*>> addNewSymbolsConsideringJvmFacades(
         destination: MutableList<S>,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/providers/LLFirProvider.kt

@@ -113,6 +113,11 @@ internal class LLFirProvider(
         override fun getPackage(fqName: FqName): FqName? =
             providerHelper.getPackage(fqName)
 
+        // TODO: Consider having proper implementations for sake of optimizations
+        override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
+        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
+        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
+
         override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
             return getFirClassifierByFqName(classId)?.symbol
         }

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(

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

@@ -124,7 +124,25 @@ abstract class AbstractFirDeserializedSymbolProvider(
     // This method should only be used for sake of optimization to avoid having too many empty-list/null values in our caches
     protected abstract fun computePackageSetWithNonClassDeclarations(): Set<String>
 
-    protected abstract fun mayHaveTopLevelClass(classId: ClassId): Boolean
+    override fun computePackageSetWithTopLevelCallables(): Set<String> = computePackageSetWithNonClassDeclarations()
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> = allNamesByPackage.getValue(packageFqName)
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? {
+        val classesInPackage = knownTopLevelClassesInPackage(packageFqName) ?: return null
+
+        if (packageFqName.asString() !in packageNamesForNonClassDeclarations) return classesInPackage
+
+        val typeAliasNames = typeAliasesNamesByPackage.getValue(packageFqName)
+        if (typeAliasNames.isEmpty()) return classesInPackage
+
+        return buildSet {
+            addAll(classesInPackage)
+            typeAliasNames.mapTo(this) { it.asString() }
+        }
+    }
+
+    protected abstract fun knownTopLevelClassesInPackage(packageFqName: FqName): Set<String>?
 
     protected abstract fun extractClassMetadata(
         classId: ClassId,
@@ -240,6 +258,11 @@ abstract class AbstractFirDeserializedSymbolProvider(
         return classCache.getValue(classId, parentContext)
     }
 
+    private fun mayHaveTopLevelClass(topLevelClassId: ClassId): Boolean {
+        val knownClassNames = knownTopLevelClassifiersInPackage(topLevelClassId.packageFqName) ?: return true
+        return topLevelClassId.shortClassName.asString() in knownClassNames
+    }
+
     private fun getTypeAlias(classId: ClassId): FirTypeAliasSymbol? {
         if (!classId.relativeClassName.isOneSegmentFQN()) return null
 

compiler/fir/fir-deserialization/src/org/jetbrains/kotlin/fir/resolve/providers/impl/FirBuiltinSymbolProvider.kt

@@ -77,6 +77,18 @@ open class FirBuiltinSymbolProvider(
         } ?: syntheticFunctionalInterfaceCache.tryGetSyntheticFunctionalInterface(classId)
     }
 
+    override fun computePackageSetWithTopLevelCallables(): Set<String> =
+        allPackageFragments.keys.mapTo(mutableSetOf()) { it.asString() }
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
+        allPackageFragments[packageFqName]?.flatMapTo(mutableSetOf()) { fragment ->
+            fragment.classDataFinder.allClassIds.map { it.shortClassName.asString() }
+        }.orEmpty()
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> =
+        allPackageFragments[packageFqName]?.flatMapTo(mutableSetOf()) {
+            it.getTopLevelCallableNames()
+        }.orEmpty()
 
     @FirSymbolProviderInternals
     override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
@@ -154,6 +166,9 @@ open class FirBuiltinSymbolProvider(
             return getTopLevelFunctionSymbols(name)
         }
 
+        fun getTopLevelCallableNames(): Collection<Name> =
+            packageProto.`package`.functionList.map { nameResolver.getName(it.name) }
+
         fun getTopLevelFunctionSymbols(name: Name): List<FirNamedFunctionSymbol> {
             return packageProto.`package`.functionList.filter { nameResolver.getName(it.name) == name }.map {
                 memberDeserializer.loadFunction(it).symbol

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

@@ -85,6 +85,8 @@ abstract class FirJavaFacade(
         return classId.relativeClassName.topLevelName() in knownNames
     }
 
+    fun knownClassNamesInPackage(packageFqName: FqName): Set<String>? = knownClassNamesInPackage.getValue(packageFqName)
+
     abstract fun getModuleDataForClass(javaClass: JavaClass): FirModuleData
 
     private fun JavaTypeParameter.toFirTypeParameter(

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

@@ -45,6 +45,12 @@ class JavaSymbolProvider(
     override fun getClassLikeSymbolByClassId(classId: ClassId): FirRegularClassSymbol? =
         if (javaFacade.hasTopLevelClassOf(classId)) getFirJavaClass(classId) else null
 
+    override fun computePackageSetWithTopLevelCallables(): Set<String> = emptySet()
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = javaFacade.knownClassNamesInPackage(packageFqName)
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> = emptySet()
+
     private fun getFirJavaClass(classId: ClassId): FirRegularClassSymbol? =
         classCache.getValue(classId, classId.outerClassId?.let { getFirJavaClass(it) })
 

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/deserialization/JvmClassFileBasedSymbolProvider.kt

@@ -99,7 +99,7 @@ class JvmClassFileBasedSymbolProvider(
 
     override fun computePackageSetWithNonClassDeclarations(): Set<String> = packagePartProvider.computePackageSetWithNonClassDeclarations()
 
-    override fun mayHaveTopLevelClass(classId: ClassId): Boolean = javaFacade.hasTopLevelClassOf(classId)
+    override fun knownTopLevelClassesInPackage(packageFqName: FqName): Set<String>? = javaFacade.knownClassNamesInPackage(packageFqName)
 
     private val KotlinJvmBinaryClass.incompatibility: IncompatibleVersionErrorData<JvmMetadataVersion>?
         get() {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/deserialization/OptionalAnnotationClassesProvider.kt

@@ -45,13 +45,22 @@ class OptionalAnnotationClassesProvider(
         return@lazy Pair(optionalAnnotationClasses, optionalAnnotationPackages)
     }
 
+    private val optionalAnnotationClassNamesByPackage: Map<FqName, Set<String>> by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        buildMap<FqName, MutableSet<String>> {
+            for (classId in optionalAnnotationClassesAndPackages.first.keys) {
+                getOrPut(classId.packageFqName, ::mutableSetOf).add(classId.shortClassName.asString())
+            }
+        }
+    }
+
     override fun computePackagePartsInfos(packageFqName: FqName): List<PackagePartsCacheData> {
         return emptyList()
     }
 
     override fun computePackageSetWithNonClassDeclarations(): Set<String> = optionalAnnotationClassesAndPackages.second
 
-    override fun mayHaveTopLevelClass(classId: ClassId): Boolean = classId in optionalAnnotationClassesAndPackages.first
+    override fun knownTopLevelClassesInPackage(packageFqName: FqName): Set<String> =
+        optionalAnnotationClassNamesByPackage[packageFqName] ?: emptySet()
 
     override fun extractClassMetadata(
         classId: ClassId,

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirExtensionDeclarationsSymbolProvider.kt

@@ -53,6 +53,35 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
         hasPackage(packageFqName)
     }
 
+    private val callableNamesInPackageCache: FirLazyValue<Map<FqName, Set<Name>>, Nothing?> =
+        cachesFactory.createLazyValue {
+            computeNamesGroupedByPackage(
+                FirDeclarationGenerationExtension::getTopLevelCallableIds,
+                CallableId::packageName, CallableId::callableName
+            )
+        }
+
+    private val classNamesInPackageCache: FirLazyValue<Map<FqName, Set<String>>, Nothing?> =
+        cachesFactory.createLazyValue {
+            computeNamesGroupedByPackage(
+                FirDeclarationGenerationExtension::getTopLevelClassIds,
+                ClassId::getPackageFqName
+            ) { it.shortClassName.asString() }
+        }
+
+    private fun <I, N> computeNamesGroupedByPackage(
+        ids: FirDeclarationGenerationExtension.() -> Collection<I>,
+        packageFqName: (I) -> FqName,
+        shortName: (I) -> N,
+    ): Map<FqName, Set<N>> =
+        buildMap<FqName, MutableSet<N>> {
+            for (extension in extensions) {
+                for (id in extension.ids()) {
+                    getOrPut(packageFqName(id)) { mutableSetOf() }.add(shortName(id))
+                }
+            }
+        }
+
     private val extensionsByTopLevelClassId: FirLazyValue<Map<ClassId, List<FirDeclarationGenerationExtension>>, Nothing?> =
         session.firCachesFactory.createLazyValue {
             extensions.flatGroupBy { it.topLevelClassIdsCache.getValue() }
@@ -139,4 +168,15 @@ class FirExtensionDeclarationsSymbolProvider private constructor(
     override fun getPackage(fqName: FqName): FqName? {
         return fqName.takeIf { packageCache.getValue(fqName, null) }
     }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String> =
+        extensions.flatMapTo(mutableSetOf()) { extension ->
+            extension.topLevelCallableIdsCache.getValue(null).map { it.packageName.asString() }
+        }
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
+        classNamesInPackageCache.getValue()[packageFqName] ?: emptySet()
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> =
+        callableNamesInPackageCache.getValue()[packageFqName].orEmpty()
 }

compiler/fir/providers/src/org/jetbrains/kotlin/fir/extensions/FirSwitchableExtensionDeclarationsSymbolProvider.kt

@@ -68,6 +68,16 @@ class FirSwitchableExtensionDeclarationsSymbolProvider private constructor(
     fun enable() {
         disabled = false
     }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? =
+        if (disabled) null else delegate.computePackageSetWithTopLevelCallables()
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? =
+        if (disabled) null else delegate.knownTopLevelClassifiersInPackage(packageFqName)
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? =
+        if (disabled) null else delegate.computeCallableNamesInPackage(packageFqName)
+
 }
 
 val FirSession.generatedDeclarationsSymbolProvider: FirSwitchableExtensionDeclarationsSymbolProvider? by FirSession.nullableSessionComponentAccessor()

compiler/fir/providers/src/org/jetbrains/kotlin/fir/resolve/providers/FirSymbolProvider.kt

@@ -55,8 +55,37 @@ abstract class FirSymbolProvider(val session: FirSession) : FirSessionComponent
     abstract fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name)
 
     abstract fun getPackage(fqName: FqName): FqName? // TODO: Replace to symbol sometime
+
+    /**
+     * All the three "compute*" functions below have the following common contract:
+     * - They return null in case necessary name set is too hard/impossible to compute.
+     * - They might return a strict superset of the name set, i.e. the resulting set might contain some names that do not belong to the provider.
+     * - It might be non-cheap to compute them on each query, thus their result should be cached properly.
+     *
+     * @returns full package names that contain some top-level callables
+     */
+    abstract fun computePackageSetWithTopLevelCallables(): Set<String>?
+
+    /**
+     * @returns top-level classifier names that belong to `packageFqName` or null if it's complicated to compute the set
+     *
+     * All usages must take into account that the result might not include kotlin.FunctionN
+     * (and others for which org.jetbrains.kotlin.builtins.functions.FunctionClassKind.Companion.byClassNamePrefix not-null)
+     */
+    abstract fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>?
+
+    /**
+     * @returns top-level callable names that belong to `packageFqName` or null if it's complicated to compute the set
+     */
+    abstract fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>?
 }
 
+/**
+ * Works almost as regular flatMap, but returns a set and returns null if any lambda call returned null
+ */
+inline fun <T, R> Iterable<T>.flatMapToNullableSet(transform: (T) -> Iterable<R>?): Set<R>? =
+    flatMapTo(mutableSetOf()) { transform(it) ?: return null }
+
 private fun FirSymbolProvider.getClassDeclaredMemberScope(classId: ClassId): FirScope? {
     val classSymbol = getClassLikeSymbolByClassId(classId) as? FirRegularClassSymbol ?: return null
     return session.declaredMemberScope(classSymbol.fir)

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

@@ -0,0 +1,151 @@
+/*
+ * Copyright 2010-2018 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.fir.resolve.providers.impl
+
+import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.NoMutableState
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.createCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.resolve.providers.flatMapToNullableSet
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+@NoMutableState
+class FirCachingCompositeSymbolProvider(
+    session: FirSession,
+    val providers: List<FirSymbolProvider>,
+    // This property is necessary just to make sure we don't use the hack at `createCopyWithCleanCaches` more than once or in cases
+    // we are not assumed to use it.
+    private val expectedCachesToBeCleanedOnce: Boolean = false,
+) : FirSymbolProvider(session) {
+
+    private val classLikeCache = session.firCachesFactory.createCache(::computeClass)
+    private val topLevelCallableCache = session.firCachesFactory.createCache(::computeTopLevelCallables)
+    private val topLevelFunctionCache = session.firCachesFactory.createCache(::computeTopLevelFunctions)
+    private val topLevelPropertyCache = session.firCachesFactory.createCache(::computeTopLevelProperties)
+    private val packageCache = session.firCachesFactory.createCache(::computePackage)
+
+    private val callablePackageSet: Set<String>? by lazy(LazyThreadSafetyMode.PUBLICATION) {
+        computePackageSetWithTopLevelCallables().also {
+            ensureNotNull(it) { "package names with callables" }
+        }
+    }
+
+    private val knownTopLevelClassifierNamesInPackage: FirCache<FqName, Set<String>?, Nothing?> =
+        session.firCachesFactory.createCache { packageFqName ->
+            knownTopLevelClassifiersInPackage(packageFqName).also {
+                ensureNotNull(it) { "classifier names in package $packageFqName" }
+            }
+        }
+
+    private val callableNamesInPackage: FirCache<FqName, Set<Name>?, Nothing?> =
+        session.firCachesFactory.createCache { packageFqName ->
+            computeCallableNamesInPackage(packageFqName).also {
+                ensureNotNull(it) { "callable names in package $packageFqName" }
+            }
+        }
+
+    private inline fun ensureNotNull(v: Any?, representation: () -> String) {
+        require(v != null || expectedCachesToBeCleanedOnce) {
+            "${representation()} is expected to be not null in CLI"
+        }
+    }
+
+    // Unfortunately, this is a part of a hack for overcoming the problem of plugin's generated entities
+    // (for more details see its usage at org.jetbrains.kotlin.fir.resolve.transformers.plugin.FirCompilerRequiredAnnotationsResolveProcessor.afterPhase)
+    fun createCopyWithCleanCaches(): FirCachingCompositeSymbolProvider {
+        require(expectedCachesToBeCleanedOnce) { "Unexpected caches clearing" }
+        return FirCachingCompositeSymbolProvider(session, providers, expectedCachesToBeCleanedOnce = false)
+    }
+
+    override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
+        if (!mayHaveTopLevelCallablesInPackage(packageFqName, name)) return emptyList()
+        return topLevelCallableCache.getValue(CallableId(packageFqName, name))
+    }
+
+    private fun mayHaveTopLevelCallablesInPackage(packageFqName: FqName, name: Name): Boolean {
+        if (callablePackageSet != null && packageFqName.asString() !in callablePackageSet!!) return false
+        val callableNamesInPackage = callableNamesInPackage.getValue(packageFqName) ?: return true
+        return name in callableNamesInPackage
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+        destination += getTopLevelCallableSymbols(packageFqName, name)
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+        if (!mayHaveTopLevelCallablesInPackage(packageFqName, name)) return
+        destination += topLevelFunctionCache.getValue(CallableId(packageFqName, name))
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+        if (!mayHaveTopLevelCallablesInPackage(packageFqName, name)) return
+        destination += topLevelPropertyCache.getValue(CallableId(packageFqName, name))
+    }
+
+    override fun getPackage(fqName: FqName): FqName? {
+        return packageCache.getValue(fqName)
+    }
+
+    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        val knownClassifierNames = knownTopLevelClassifierNamesInPackage.getValue(classId.packageFqName)
+        if (knownClassifierNames != null && !isNameForFunctionClass(classId)) {
+            val outerClassId = classId.outerClassId
+            if (outerClassId == null && classId.shortClassName.asString() !in knownClassifierNames) return null
+            if (outerClassId != null && classId.outermostClassId.shortClassName.asString() !in knownClassifierNames) return null
+        }
+
+        return classLikeCache.getValue(classId)
+    }
+
+    private fun isNameForFunctionClass(classId: ClassId): Boolean {
+        return FunctionClassKind.byClassNamePrefix(classId.packageFqName, classId.shortClassName.asString()) != null
+    }
+
+    @OptIn(FirSymbolProviderInternals::class)
+    private fun computeTopLevelCallables(callableId: CallableId): List<FirCallableSymbol<*>> = buildList {
+        providers.forEach { it.getTopLevelCallableSymbolsTo(this, callableId.packageName, callableId.callableName) }
+    }
+
+    @OptIn(FirSymbolProviderInternals::class)
+    private fun computeTopLevelFunctions(callableId: CallableId): List<FirNamedFunctionSymbol> = buildList {
+        providers.forEach { it.getTopLevelFunctionSymbolsTo(this, callableId.packageName, callableId.callableName) }
+    }
+
+    @OptIn(FirSymbolProviderInternals::class)
+    private fun computeTopLevelProperties(callableId: CallableId): List<FirPropertySymbol> = buildList {
+        providers.forEach { it.getTopLevelPropertySymbolsTo(this, callableId.packageName, callableId.callableName) }
+    }
+
+    private fun computePackage(it: FqName): FqName? =
+        providers.firstNotNullOfOrNull { provider -> provider.getPackage(it) }
+
+    private fun computeClass(classId: ClassId): FirClassLikeSymbol<*>? =
+        providers.firstNotNullOfOrNull { provider -> provider.getClassLikeSymbolByClassId(classId) }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? =
+        providers.flatMapToNullableSet { it.computePackageSetWithTopLevelCallables() }
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? =
+        providers.flatMapToNullableSet { it.knownTopLevelClassifiersInPackage(packageFqName) }
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? =
+        providers.flatMapToNullableSet { it.computeCallableNamesInPackage(packageFqName) }
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2023 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.
  */
 
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.NoMutableState
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.resolve.providers.flatMapToNullableSet
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
@@ -49,4 +50,13 @@ class FirCompositeSymbolProvider(session: FirSession, val providers: List<FirSym
     override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
         return providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classId) }
     }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? =
+        providers.flatMapToNullableSet { it.computePackageSetWithTopLevelCallables() }
+
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? =
+        providers.flatMapToNullableSet { it.knownTopLevelClassifiersInPackage(packageFqName) }
+
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? =
+        providers.flatMapToNullableSet { it.computeCallableNamesInPackage(packageFqName) }
 }

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