[LL] [cls] don't access index to retrieve fir for given ktElement

also don't check name pattern, because given ktElement proves that corresponding package contains something
2023-05-04 23:45:56 +02:00
parent e47561e70b
commit cf6a06fd14
5 changed files with 93 additions and 26 deletions
@@ -5,6 +5,7 @@

 package org.jetbrains.kotlin.analysis.low.level.api.fir.providers

+import org.jetbrains.kotlin.analysis.low.level.api.fir.stubBased.deserialization.JvmStubBasedFirDeserializedSymbolProvider
 import org.jetbrains.kotlin.analysis.utils.collections.buildSmartList
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
@@ -14,12 +15,14 @@ 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.load.kotlin.FacadeClassSource
-import org.jetbrains.kotlin.load.kotlin.JvmPackagePartSource
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtCallableDeclaration
+import org.jetbrains.kotlin.psi.KtClassLikeDeclaration
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.utils.SmartSet
+import org.jetbrains.kotlin.utils.addIfNotNull

 internal class LLFirModuleWithDependenciesSymbolProvider(
    session: FirSession,
@@ -33,20 +36,29 @@ internal class LLFirModuleWithDependenciesSymbolProvider(
    fun getClassLikeSymbolByFqNameWithoutDependencies(classId: ClassId): FirClassLikeSymbol<*>? =
        providers.firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classId) }

+    fun getClassLikeSymbolByFqNameWithoutDependencies(
+        classLikeDeclaration: KtClassLikeDeclaration,
+        classId: ClassId
+    ): FirClassLikeSymbol<*>? {
+        return providers.filterIsInstance(JvmStubBasedFirDeserializedSymbolProvider::class.java)
+            .firstNotNullOfOrNull { it.getClassLikeSymbolByClassId(classLikeDeclaration, classId) }
+    }
+
    @FirSymbolProviderInternals
    override fun getTopLevelCallableSymbolsTo(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-        getTopLevelCallableSymbolsToWithoutDependencies(destination, packageFqName, name)
+        providers.forEach { it.getTopLevelCallableSymbolsTo(destination, packageFqName, name) }
        dependencyProvider.getTopLevelCallableSymbolsTo(destination, packageFqName, name)
    }

    @FirSymbolProviderInternals
-    fun getTopLevelCallableSymbolsWithoutDependencies(packageFqName: FqName, name: Name): List<FirCallableSymbol<*>> {
-        return buildList { getTopLevelCallableSymbolsToWithoutDependencies(this, packageFqName, name) }
-    }
-
-    @FirSymbolProviderInternals
-    fun getTopLevelCallableSymbolsToWithoutDependencies(destination: MutableList<FirCallableSymbol<*>>, packageFqName: FqName, name: Name) {
-        providers.forEach { it.getTopLevelCallableSymbolsTo(destination, packageFqName, name) }
+    fun getTopLevelDeserializedCallableSymbolsToWithoutDependencies(
+        destination: MutableList<FirCallableSymbol<*>>,
+        callableDeclaration: KtCallableDeclaration,
+        packageFqName: FqName,
+        shortName: Name
+    ) {
+        providers.filterIsInstance(JvmStubBasedFirDeserializedSymbolProvider::class.java)
+            .forEach { destination.addIfNotNull(it.getTopLevelCallableSymbol(callableDeclaration, packageFqName, shortName)) }
    }

    @FirSymbolProviderInternals
@@ -18,15 +18,13 @@ import org.jetbrains.kotlin.fir.caches.getValue
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.deserialization.SingleModuleDataProvider
 import org.jetbrains.kotlin.fir.java.deserialization.KotlinBuiltins
+import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProviderInternals
 import org.jetbrains.kotlin.fir.scopes.FirKotlinScopeProvider
 import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.name.*
-import org.jetbrains.kotlin.psi.KtClassOrObject
-import org.jetbrains.kotlin.psi.KtNamedFunction
-import org.jetbrains.kotlin.psi.KtProperty
-import org.jetbrains.kotlin.psi.KtTypeAlias
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import org.jetbrains.kotlin.serialization.deserialization.MetadataPackageFragment

@@ -67,7 +65,7 @@ internal open class JvmStubBasedFirDeserializedSymbolProvider(

    private val typeAliasCache: FirCache<ClassId, FirTypeAliasSymbol?, StubBasedFirDeserializationContext?> =
        session.firCachesFactory.createCacheWithPostCompute(
-            createValue = { classId, _ -> findAndDeserializeTypeAlias(classId) },
+            createValue = { classId, context -> findAndDeserializeTypeAlias(classId, context) },
            postCompute = { _, symbol, postProcessor ->
                if (postProcessor != null && symbol != null) {
                    postProcessor.invoke(symbol)
@@ -93,12 +91,16 @@ internal open class JvmStubBasedFirDeserializedSymbolProvider(
        return namesByPackageCache.getTopLevelClassifierNamesInPackage(packageFqName)
    }

-    private fun findAndDeserializeTypeAlias(classId: ClassId): Pair<FirTypeAliasSymbol?, DeserializedTypeAliasPostProcessor?> {
-        val classLikeDeclaration = declarationProvider.getClassLikeDeclarationByClassId(classId)?.originalElement
+    private fun findAndDeserializeTypeAlias(
+        classId: ClassId,
+        context: StubBasedFirDeserializationContext?
+    ): Pair<FirTypeAliasSymbol?, DeserializedTypeAliasPostProcessor?> {
+        val classLikeDeclaration =
+            context?.classLikeDeclaration ?: declarationProvider.getClassLikeDeclarationByClassId(classId)?.originalElement
        if (classLikeDeclaration is KtTypeAlias) {
            val symbol = FirTypeAliasSymbol(classId)
            val postProcessor: DeserializedTypeAliasPostProcessor = {
-                val rootContext = StubBasedFirDeserializationContext.createRootContext(
+                val rootContext = context ?: StubBasedFirDeserializationContext.createRootContext(
                    moduleData,
                    StubBasedAnnotationDeserializer(session),
                    classId.packageFqName,
@@ -117,7 +119,12 @@ internal open class JvmStubBasedFirDeserializedSymbolProvider(
        classId: ClassId,
        parentContext: StubBasedFirDeserializationContext? = null
    ): FirRegularClassSymbol? {
-        val classLikeDeclaration = declarationProvider.getClassLikeDeclarationByClassId(classId)?.originalElement ?: return null
+        val (classLikeDeclaration, context) =
+            if (parentContext?.classLikeDeclaration != null) {
+                parentContext.classLikeDeclaration to null
+            } else {
+                (declarationProvider.getClassLikeDeclarationByClassId(classId)?.originalElement ?: return null) to parentContext
+            }
        val symbol = FirRegularClassSymbol(classId)
        if (classLikeDeclaration is KtClassOrObject) {
            deserializeClassToSymbol(
@@ -128,14 +135,14 @@ internal open class JvmStubBasedFirDeserializedSymbolProvider(
                moduleData,
                StubBasedAnnotationDeserializer(session),
                kotlinScopeProvider,
-                parentContext,
+                parentContext = context,
                containerSource = if (initialOrigin == FirDeclarationOrigin.BuiltIns) null else JvmFromStubDecompilerSource(
                    JvmClassName.byClassId(
                        classId
                    )
                ),
                deserializeNestedClass = this::getClass,
-                initialOrigin
+                initialOrigin = initialOrigin
            )
            return symbol
        }
@@ -223,4 +230,51 @@ internal open class JvmStubBasedFirDeserializedSymbolProvider(
        if (!namesByPackageCache.mayHaveTopLevelClassifier(classId, mayHaveFunctionClass = false)) return null
        return getClass(classId) ?: getTypeAlias(classId)
    }
+
+    fun getClassLikeSymbolByClassId(classLikeDeclaration: KtClassLikeDeclaration, classId: ClassId): FirClassLikeSymbol<*>? {
+        val annotationDeserializer = StubBasedAnnotationDeserializer(session)
+        val deserializationContext = StubBasedFirDeserializationContext(
+            moduleData,
+            classId.packageFqName,
+            classId.relativeClassName,
+            StubBasedFirTypeDeserializer(
+                moduleData,
+                annotationDeserializer,
+                parent = null,
+                containingSymbol = null,
+                owner = null,
+                initialOrigin
+            ),
+            annotationDeserializer,
+            containerSource = null,
+            outerClassSymbol = null,
+            outerTypeParameters = emptyList(),
+            initialOrigin,
+            classLikeDeclaration
+        )
+        if (classLikeDeclaration is KtClassOrObject) {
+            return classCache.getValue(
+                classId,
+                deserializationContext
+            )
+        }
+        return typeAliasCache.getValue(classId, deserializationContext)
+    }
+
+    fun getTopLevelCallableSymbol(
+        callableDeclaration: KtCallableDeclaration,
+        packageFqName: FqName,
+        shortName: Name
+    ): FirCallableSymbol<*>? {
+        //possible overloads spoils here
+        //we can't use only this callable instead of index access to fill the cache
+        //names check is redundant though as we already have existing callable in scope
+        val callableId = CallableId(packageFqName, shortName)
+        val callableSymbols = when (callableDeclaration) {
+            is KtNamedFunction -> functionCache.getValue(callableId)
+            is KtProperty -> propertyCache.getValue(callableId)
+            else -> null
+        }
+        return callableSymbols?.singleOrNull { it.fir.realPsi == callableDeclaration }
+    }
 }
@@ -45,7 +45,8 @@ internal class StubBasedFirDeserializationContext(
    val containerSource: DeserializedContainerSource?,
    val outerClassSymbol: FirRegularClassSymbol?,
    val outerTypeParameters: List<FirTypeParameterSymbol>,
-    private val initialOrigin: FirDeclarationOrigin
+    private val initialOrigin: FirDeclarationOrigin,
+    val classLikeDeclaration: KtClassLikeDeclaration? = null
 ) {
    val session: FirSession = moduleData.session

@@ -42,7 +42,7 @@ internal class StubBasedFirTypeDeserializer(
    private val annotationDeserializer: StubBasedAnnotationDeserializer,
    private val parent: StubBasedFirTypeDeserializer?,
    private val containingSymbol: FirBasedSymbol<*>?,
-    owner: KtTypeParameterListOwner,
+    owner: KtTypeParameterListOwner?,
    initialOrigin: FirDeclarationOrigin
 ) {
    private val typeParametersByName: Map<String, FirTypeParameterSymbol>
@@ -51,8 +51,8 @@ internal class StubBasedFirTypeDeserializer(
        get() = typeParametersByName.values.toList()

    init {
-        val typeParameters = owner.typeParameters
-        if (typeParameters.isNotEmpty()) {
+        val typeParameters = owner?.typeParameters
+        if (!typeParameters.isNullOrEmpty()) {
            typeParametersByName = mutableMapOf()
            val builders = mutableListOf<FirTypeParameterBuilder>()
            for (typeParameter in typeParameters) {
@@ -70,7 +70,7 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid

        val classCandidate = when (symbolProvider) {
            is LLFirModuleWithDependenciesSymbolProvider -> {
-                symbolProvider.getClassLikeSymbolByFqNameWithoutDependencies(classId)
+                symbolProvider.getClassLikeSymbolByFqNameWithoutDependencies(declaration, classId)
                    ?: symbolProvider.friendBuiltinsProvider?.getClassLikeSymbolByClassId(classId)
            }
            else -> {
@@ -162,7 +162,7 @@ private fun FirSymbolProvider.findCallableCandidates(
        @OptIn(FirSymbolProviderInternals::class)
        return when (this) {
            is LLFirModuleWithDependenciesSymbolProvider -> buildList {
-                getTopLevelCallableSymbolsToWithoutDependencies(this, packageFqName, shortName)
+                getTopLevelDeserializedCallableSymbolsToWithoutDependencies(this, declaration, packageFqName, shortName)
                friendBuiltinsProvider?.getTopLevelCallableSymbolsTo(this, packageFqName, shortName)
            }
            else -> getTopLevelCallableSymbols(packageFqName, shortName)