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%)
2023-01-06 17:29:44 +01:00
parent 6705d211a6
commit 9c988fd8d8
18 changed files with 437 additions and 44 deletions
@@ -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()
 }
@@ -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()
@@ -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)
@@ -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) }
+}
@@ -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) }
 }