[LL FIR] KTIJ-24296 Extract symbol provider name cache into a class

- The caching logic from `LLFirProvider` may be used in other symbol providers.
2023-02-16 14:43:20 +01:00
parent d9a024bab2
commit ece038b49b
5 changed files with 109 additions and 74 deletions
@@ -109,60 +109,53 @@ internal class LLFirProvider(
    @NoMutableState
    private inner class SymbolProvider : FirSymbolProvider(session) {
        override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
-            val classifierNamesInPackage = providerHelper.getTopLevelClassifierNamesInPackage(classId.packageFqName)
-            if (!mayHaveTopLevelClassifier(classId, classifierNamesInPackage.names, classifierNamesInPackage.mayHaveFunctionClass)) {
-                return null
-            }
-
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelClassifier(classId)) return null
            return getFirClassifierByFqName(classId)?.symbol
        }

        override fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
            return providerHelper.getTopLevelCallableSymbols(packageFqName, name)
        }

        @FirSymbolProviderInternals
        override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return
            destination += providerHelper.getTopLevelCallableSymbols(packageFqName, name)
        }

        override fun getTopLevelFunctionSymbols(packageFqName: FqName, name: Name): List<FirNamedFunctionSymbol> {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
            return providerHelper.getTopLevelFunctionSymbols(packageFqName, name)
        }

        @FirSymbolProviderInternals
        override fun getTopLevelFunctionSymbolsTo(destination: MutableList<FirNamedFunctionSymbol>, packageFqName: FqName, name: Name) {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return
            destination += providerHelper.getTopLevelFunctionSymbols(packageFqName, name)
        }

        override fun getTopLevelPropertySymbols(packageFqName: FqName, name: Name): List<FirPropertySymbol> {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return emptyList()
            return providerHelper.getTopLevelPropertySymbols(packageFqName, name)
        }

        @FirSymbolProviderInternals
        override fun getTopLevelPropertySymbolsTo(destination: MutableList<FirPropertySymbol>, packageFqName: FqName, name: Name) {
-            if (!mayHaveTopLevelCallable(packageFqName, name)) return
+            if (!providerHelper.symbolNameCache.mayHaveTopLevelCallable(packageFqName, name)) return
            destination += providerHelper.getTopLevelPropertySymbols(packageFqName, name)
        }

-        private fun mayHaveTopLevelCallable(packageFqName: FqName, name: Name): Boolean =
-            name in computeCallableNamesInPackage(packageFqName)
-
        override fun getPackage(fqName: FqName): FqName? =
            providerHelper.getPackage(fqName)

        // Computing the set of such package names is expensive and would require a new index. For now, it is not worth the marginal gains.
        override fun computePackageSetWithTopLevelCallables(): Set<String>? = null

-        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String> =
-            providerHelper.getTopLevelClassifierNamesInPackage(packageFqName).names
+        override fun knownTopLevelClassifiersInPackage(packageFqName: FqName): Set<String>? =
+            providerHelper.symbolNameCache.getTopLevelClassifierNamesInPackage(packageFqName)?.names

-        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name> =
-            providerHelper.getTopLevelCallableNamesInPackage(packageFqName)
+        override fun computeCallableNamesInPackage(packageFqName: FqName): Set<Name>? =
+            providerHelper.symbolNameCache.getTopLevelCallableNamesInPackage(packageFqName)
    }
 }
@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.providers

 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirSymbolProviderNameCache
 import org.jetbrains.kotlin.analysis.providers.KotlinDeclarationProvider
 import org.jetbrains.kotlin.analysis.providers.KotlinPackageProvider
 import org.jetbrains.kotlin.builtins.StandardNames
@@ -22,7 +23,6 @@ import org.jetbrains.kotlin.fir.languageVersionSettings
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.functionTypeService
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -49,26 +49,6 @@ internal class LLFirProviderHelper(
                ?: error("Classifier $classId was found in file ${ktClass.containingKtFile.virtualFilePath} but was not found in FirFile")
        }

-    internal class ClassifierNamesInPackage(
-        val names: Set<String>,
-
-        /**
-         * Some packages (such as `kotlin` or `kotlin.coroutines`) may contain function classes generated by the compiler. To avoid a hash
-         * map access when checking `mayHaveTopLevelClassifier`, [mayHaveFunctionClass] is precomputed for each package name. Also see
-         * [org.jetbrains.kotlin.builtins.functions.FunctionTypeKind].
-         */
-        val mayHaveFunctionClass: Boolean,
-    )
-
-    private val topLevelClassifierNamesByPackage =
-        firSession.firCachesFactory.createCache<FqName, ClassifierNamesInPackage> { packageFqName ->
-            val names = declarationProvider
-                .getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName)
-                .mapTo(mutableSetOf()) { it.asString() }
-            val mayHaveFunctionClass = firSession.functionTypeService.hasKindWithSpecificPackage(packageFqName)
-            ClassifierNamesInPackage(names, mayHaveFunctionClass)
-        }
-
    private val callablesByCallableId = firSession.firCachesFactory.createCache<CallableId, List<FirCallableSymbol<*>>> { callableId ->
        val files = declarationProvider.getTopLevelCallableFiles(callableId).ifEmpty { return@createCache emptyList() }
        buildList {
@@ -79,10 +59,15 @@ internal class LLFirProviderHelper(
        }
    }

-    private val topLevelCallableNamesByPackage =
-        firSession.firCachesFactory.createCache<FqName, Set<Name>> { packageFqName ->
+    val symbolNameCache = object : LLFirSymbolProviderNameCache(firSession) {
+        override fun computeClassifierNames(packageFqName: FqName): Set<String>? =
+            declarationProvider
+                .getTopLevelKotlinClassLikeDeclarationNamesInPackage(packageFqName)
+                .mapTo(mutableSetOf()) { it.asString() }
+
+        override fun computeCallableNames(packageFqName: FqName): Set<Name>? =
            declarationProvider.getTopLevelCallableNamesInPackage(packageFqName)
-        }
+    }

    fun getFirClassifierByFqNameAndDeclaration(
        classId: ClassId,
@@ -93,9 +78,6 @@ internal class LLFirProviderHelper(
        return classifierByClassId.getValue(classId, classLikeDeclaration)
    }

-    fun getTopLevelClassifierNamesInPackage(packageFqName: FqName): ClassifierNamesInPackage =
-        topLevelClassifierNamesByPackage.getValue(packageFqName)
-
    fun getTopLevelCallableSymbols(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
        if (!allowKotlinPackage && packageFqName.isKotlinPackage()) return emptyList()
        val callableId = CallableId(packageFqName, name)
@@ -110,9 +92,6 @@ internal class LLFirProviderHelper(
        return getTopLevelCallableSymbols(packageFqName, name).filterIsInstance<FirPropertySymbol>()
    }

-    fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name> =
-        topLevelCallableNamesByPackage.getValue(packageFqName)
-
    private fun FirFile.collectCallableDeclarationsTo(list: MutableList<FirCallableSymbol<*>>, name: Name) {
        declarations.mapNotNullTo(list) { declaration ->
            if (declaration is FirCallableDeclaration && declaration.symbol.callableId.callableName == name) {
@@ -0,0 +1,62 @@
+/*
+ * 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.util
+
+import org.jetbrains.kotlin.fir.FirSession
+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.mayHaveTopLevelClassifier
+import org.jetbrains.kotlin.fir.types.functionTypeService
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+/**
+ * Caches the names of classifiers and callables contained in a package. [LLFirSymbolProviderNameCache] is used by symbol providers to abort
+ * symbol finding early if the symbol name isn't contained in the symbol provider's domain.
+ */
+internal abstract class LLFirSymbolProviderNameCache(private val firSession: FirSession) {
+    internal data class ClassifierNamesInPackage(
+        val names: Set<String>,
+
+        /**
+         * Some packages (such as `kotlin` or `kotlin.coroutines`) may contain function classes generated by the compiler. To avoid a hash
+         * map access when checking `mayHaveTopLevelClassifier`, [mayHaveFunctionClass] is precomputed for each package name. Also see
+         * [org.jetbrains.kotlin.builtins.functions.FunctionTypeKind].
+         */
+        val mayHaveFunctionClass: Boolean,
+    )
+
+    abstract fun computeClassifierNames(packageFqName: FqName): Set<String>?
+    abstract fun computeCallableNames(packageFqName: FqName): Set<Name>?
+
+    private val topLevelClassifierNamesByPackage =
+        firSession.firCachesFactory.createCache<FqName, ClassifierNamesInPackage?> { packageFqName ->
+            val names = computeClassifierNames(packageFqName) ?: return@createCache null
+            val mayHaveFunctionClass = firSession.functionTypeService.hasKindWithSpecificPackage(packageFqName)
+            ClassifierNamesInPackage(names, mayHaveFunctionClass)
+        }
+
+    private val topLevelCallableNamesByPackage =
+        firSession.firCachesFactory.createCache<FqName, Set<Name>?>(::computeCallableNames)
+
+    fun getTopLevelClassifierNamesInPackage(packageFqName: FqName): ClassifierNamesInPackage? =
+        topLevelClassifierNamesByPackage.getValue(packageFqName)
+
+    fun getTopLevelCallableNamesInPackage(packageFqName: FqName): Set<Name>? =
+        topLevelCallableNamesByPackage.getValue(packageFqName)
+
+    fun mayHaveTopLevelClassifier(classId: ClassId): Boolean {
+        val (names, mayHaveFunctionClass) = getTopLevelClassifierNamesInPackage(classId.packageFqName) ?: return true
+        return names.mayHaveTopLevelClassifier(classId, firSession, mayHaveFunctionClass)
+    }
+
+    fun mayHaveTopLevelCallable(packageFqName: FqName, name: Name): Boolean {
+        val names = getTopLevelCallableNamesInPackage(packageFqName) ?: return true
+        return name in names
+    }
+}
@@ -79,31 +79,6 @@ abstract class FirSymbolProvider(val session: FirSession) : FirSessionComponent
     * @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>?
-
-    /**
-     * Whether [classId] may be contained in the [knownClassifierNames].
-     *
-     * If it's certain that [classId] cannot be a function class (for example when [classId] is known to come from a package without
-     * function classes), [mayBeFunctionClass] can be set to `false`. This avoids a hash map access in
-     * [org.jetbrains.kotlin.builtins.functions.FunctionTypeKindExtractor.getFunctionalClassKindWithArity].
-     */
-    protected fun mayHaveTopLevelClassifier(
-        classId: ClassId,
-        knownClassifierNames: Set<String>,
-        mayBeFunctionClass: Boolean = true,
-    ): Boolean {
-        if (mayBeFunctionClass && isNameForFunctionClass(classId)) return true
-
-        val outerClassId = classId.outerClassId
-        if (outerClassId == null && classId.shortClassName.asString() !in knownClassifierNames) return false
-        if (outerClassId != null && classId.outermostClassId.shortClassName.asString() !in knownClassifierNames) return false
-
-        return true
-    }
-
-    private fun isNameForFunctionClass(classId: ClassId): Boolean {
-        return session.functionTypeService.getKindByClassNamePrefix(classId.packageFqName, classId.shortClassName.asString()) != null
-    }
 }

 /**
@@ -142,6 +117,31 @@ fun FirSymbolProvider.getRegularClassSymbolByClassId(classId: ClassId): FirRegul
    return getClassLikeSymbolByClassId(classId) as? FirRegularClassSymbol
 }

+/**
+ * Whether [classId] may be contained in the set of known classifier names.
+ *
+ * If it's certain that [classId] cannot be a function class (for example when [classId] is known to come from a package without
+ * function classes), [mayBeFunctionClass] can be set to `false`. This avoids a hash map access in
+ * [org.jetbrains.kotlin.builtins.functions.FunctionTypeKindExtractor.getFunctionalClassKindWithArity].
+ */
+fun Set<String>.mayHaveTopLevelClassifier(
+    classId: ClassId,
+    session: FirSession,
+    mayBeFunctionClass: Boolean = true,
+): Boolean {
+    if (mayBeFunctionClass && isNameForFunctionClass(classId, session)) return true
+
+    val outerClassId = classId.outerClassId
+    if (outerClassId == null && classId.shortClassName.asString() !in this) return false
+    if (outerClassId != null && classId.outermostClassId.shortClassName.asString() !in this) return false
+
+    return true
+}
+
+private fun isNameForFunctionClass(classId: ClassId, session: FirSession): Boolean {
+    return session.functionTypeService.getKindByClassNamePrefix(classId.packageFqName, classId.shortClassName.asString()) != null
+}
+
 fun ClassId.toSymbol(session: FirSession): FirClassifierSymbol<*>? {
    return session.symbolProvider.getClassLikeSymbolByClassId(this)
 }
@@ -14,6 +14,7 @@ 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.resolve.providers.mayHaveTopLevelClassifier
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassLikeSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
@@ -106,7 +107,7 @@ class FirCachingCompositeSymbolProvider(

    override fun getClassLikeSymbolByClassId(classId: ClassId): FirClassLikeSymbol<*>? {
        val knownClassifierNames = knownTopLevelClassifierNamesInPackage.getValue(classId.packageFqName)
-        if (knownClassifierNames != null && !mayHaveTopLevelClassifier(classId, knownClassifierNames)) return null
+        if (knownClassifierNames != null && !knownClassifierNames.mayHaveTopLevelClassifier(classId, session)) return null

        return classLikeCache.getValue(classId)
    }