[LL FIR] KT-57220 Optimize & combine synthetic function symbol providers

- Dependency symbol providers contained a lot of useless
  `FirExtensionSyntheticFunctionInterfaceProvider`s, because such a
  provider only generates function classes from compiler plugins. If
  there aren't any such `FunctionTypeKind`s in the session, there is no
  need to even add the provider.
- `FirExtensionSyntheticFunctionInterfaceProvider`s that are added to
  the list of dependency providers despite this optimization will now be
  combined in `LLFirCombinedSyntheticFunctionSymbolProvider`. It checks
  `ClassId` heuristics once, which is more efficient than checking them
  in every provider separately.

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

@@ -0,0 +1,61 @@
+/*
+ * 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.
+ */
+
+package org.jetbrains.kotlin.analysis.low.level.api.fir.providers
+
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirSyntheticFunctionInterfaceProviderBase
+import org.jetbrains.kotlin.fir.resolve.providers.impl.FirSyntheticFunctionInterfaceProviderBase.Companion.mayBeSyntheticFunctionClassName
+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.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/**
+ * [LLFirCombinedSyntheticFunctionSymbolProvider] combines multiple synthetic function symbol providers with the advantage that [ClassId]
+ * heuristics are checked only once.
+ */
+@OptIn(FirSymbolProviderInternals::class)
+internal class LLFirCombinedSyntheticFunctionSymbolProvider(
+    session: FirSession,
+    private val providers: List<FirSyntheticFunctionInterfaceProviderBase>,
+) : FirSymbolProvider(session) {
+    private val combinedPackageNames: Set<FqName> = providers.flatMapTo(mutableSetOf()) { it.getFunctionKindPackageNames() }
+
+    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
+        if (!classId.mayBeSyntheticFunctionClassName()) return null
+        if (classId.packageFqName !in combinedPackageNames) return null
+
+        return providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassIdWithoutClassIdChecks(classId) }
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
+    }
+
+    @FirSymbolProviderInternals
+    override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
+    }
+
+    override fun getPackage(fqName: FqName): FqName? = fqName.takeIf { it in combinedPackageNames }
+
+    override fun computePackageSetWithTopLevelCallables(): Set<String>? = null
+    override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? = null
+    override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? = null
+
+    companion object {
+        fun merge(session: FirSession, providers: List<FirSyntheticFunctionInterfaceProviderBase>): FirSymbolProvider =
+            providers.singleOrNull() ?: LLFirCombinedSyntheticFunctionSymbolProvider(session, providers)
+    }
+}

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirSessionCache.kt

@@ -38,7 +38,6 @@ import org.jetbrains.kotlin.fir.extensions.*
 import org.jetbrains.kotlin.fir.java.JavaSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.*
 import org.jetbrains.kotlin.fir.resolve.providers.impl.FirExtensionSyntheticFunctionInterfaceProvider
-import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.scopes.wrapScopeWithJvmMapped
 import org.jetbrains.kotlin.fir.resolve.transformers.FirDummyCompilerLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
@@ -49,6 +48,7 @@ import org.jetbrains.kotlin.platform.jvm.JvmPlatforms
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.jvm.modules.JavaModuleResolver
 import org.jetbrains.kotlin.resolve.jvm.platform.JvmPlatformAnalyzerServices
+import org.jetbrains.kotlin.utils.addToStdlib.partitionIsInstance
 import java.util.concurrent.ConcurrentMap
 
 @OptIn(PrivateSessionConstructor::class, SessionConfiguration::class)
@@ -164,7 +164,13 @@ internal class LLFirSessionCache(private val project: Project) {
             })
 
             val javaSymbolProvider = createJavaSymbolProvider(this, moduleData, project, contentScope)
-            val syntheticFunctionalInterfaceProvider = FirExtensionSyntheticFunctionInterfaceProvider(this, moduleData, scopeProvider)
+
+            // We only need to add an extension synthetic function provider if the session's function type service even has extension kinds.
+            // Otherwise, the provider will be completely useless.
+            val syntheticFunctionalInterfaceProvider = if (this.functionTypeService.hasExtensionKinds()) {
+                FirExtensionSyntheticFunctionInterfaceProvider(this, moduleData, scopeProvider)
+            } else null
+
             register(
                 FirSymbolProvider::class,
                 LLFirModuleWithDependenciesSymbolProvider(
@@ -530,12 +536,19 @@ internal class LLFirSessionCache(private val project: Project) {
      * [session] should be the session of the dependent module. Because all symbol providers are tied to a session, we need a session to
      * create a combined symbol provider.
      */
-    @Suppress("UNUSED_PARAMETER")
     private fun List<FirSymbolProvider>.mergeDependencySymbolProvidersInto(
         session: FirSession,
         destination: MutableList<FirSymbolProvider>,
     ) {
-        destination.addAll(this)
+        val (syntheticFunctionSymbolProviders, remainingSymbolProviders1) =
+            partitionIsInstance<_, FirExtensionSyntheticFunctionInterfaceProvider>()
+
+        destination.addAll(remainingSymbolProviders1)
+
+        // Unfortunately, the functions that an extension synthetic function symbol provider might provide differ between sessions because
+        // they depend on compiler plugins. However, only extension providers that are affected by compiler plugins are added in
+        // `createSourcesSession`. We can still combine these, because the `ClassId` heuristics only need to be checked once.
+        destination.add(LLFirCombinedSyntheticFunctionSymbolProvider.merge(session, syntheticFunctionSymbolProviders))
     }
 }
 

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

@@ -69,11 +69,15 @@ abstract class FirSyntheticFunctionInterfaceProviderBase(
     val moduleData: FirModuleData,
     val kotlinScopeProvider: FirKotlinScopeProvider
 ) : FirSymbolProvider(session) {
+    @OptIn(FirSymbolProviderInternals::class)
     override fun getClassLikeSymbolByClassId(classId: ClassId): FirRegularClassSymbol? {
         if (!classId.mayBeSyntheticFunctionClassName()) return null
-        return cache.getValue(classId)
+        return getClassLikeSymbolByClassIdWithoutClassIdChecks(classId)
     }
 
+    @FirSymbolProviderInternals
+    fun getClassLikeSymbolByClassIdWithoutClassIdChecks(classId: ClassId): FirRegularClassSymbol? = cache.getValue(classId)
+
     @FirSymbolProviderInternals
     override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
     }
@@ -86,6 +90,9 @@ abstract class FirSyntheticFunctionInterfaceProviderBase(
     override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
     }
 
+    @FirSymbolProviderInternals
+    fun getFunctionKindPackageNames(): Set<FqName> = session.functionTypeService.getFunctionKindPackageNames()
+
     override fun getPackage(fqName: FqName): FqName? {
         return fqName.takeIf { session.functionTypeService.hasKindWithSpecificPackage(it) }
     }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirFunctionTypeKindService.kt

@@ -23,6 +23,16 @@ abstract class FirFunctionTypeKindService : FirSessionComponent {
         return extractor.hasKindWithSpecificPackage(packageFqName)
     }
 
+    /**
+     * Returns all package names for which [getKindByClassNamePrefix] may return a [FunctionTypeKind].
+     */
+    fun getFunctionKindPackageNames(): Set<FqName> = extractor.getFunctionKindPackageNames()
+
+    /**
+     * Whether [getKindByClassNamePrefix] may return a [FunctionTypeKind] added by a compiler plugin.
+     */
+    fun hasExtensionKinds(): Boolean = extractor.hasExtensionKinds()
+
     abstract fun extractSingleSpecialKindForFunction(functionSymbol: FirFunctionSymbol<*>): FunctionTypeKind?
     abstract fun extractAllSpecialKindsForFunction(functionSymbol: FirFunctionSymbol<*>): List<FunctionTypeKind>
     abstract fun extractAllSpecialKindsForFunctionTypeRef(typeRef: FirFunctionTypeRef): List<FunctionTypeKind>

core/compiler.common/src/org/jetbrains/kotlin/builtins/functions/FunctionTypeKindExtractor.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.name.FqName
 @RequiresOptIn
 annotation class AllowedToUsedOnlyInK1
 
-class FunctionTypeKindExtractor(kinds: List<FunctionTypeKind>) {
+class FunctionTypeKindExtractor(private val kinds: List<FunctionTypeKind>) {
     companion object {
         /**
          * This instance should be used only in:
@@ -49,6 +49,10 @@ class FunctionTypeKindExtractor(kinds: List<FunctionTypeKind>) {
         return packageFqName in knownKindsByPackageFqName
     }
 
+    fun getFunctionKindPackageNames(): Set<FqName> = knownKindsByPackageFqName.keys
+
+    fun hasExtensionKinds(): Boolean = kinds.any { !it.isBuiltin }
+
     data class KindWithArity(val kind: FunctionTypeKind, val arity: Int)
 
     private fun toInt(s: String): Int? {