[LL FIR] FirDeclarationForCompiledElementSearcher: avoid providers search for member declarations

This should improve performance and increase stability

^KT-61331
^KT-61014

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/FirDeclarationForCompiledElementSearcher.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2021 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.
  */
 
@@ -18,6 +18,9 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.*
+import org.jetbrains.kotlin.fir.scopes.getFunctions
+import org.jetbrains.kotlin.fir.scopes.getProperties
+import org.jetbrains.kotlin.fir.scopes.impl.declaredMemberScope
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
@@ -26,6 +29,7 @@ import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
 
 /**
  * Allows to search for FIR declarations by compiled [KtDeclaration]s.
@@ -36,18 +40,55 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         ProjectStructureProvider.getInstance(project)
     }
 
-    fun findNonLocalDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
-        return when (ktDeclaration) {
-            is KtEnumEntry -> findNonLocalEnumEntry(ktDeclaration)
-            is KtClassLikeDeclaration -> findNonLocalClassLikeDeclaration(ktDeclaration)
-            is KtConstructor<*> -> findConstructorOfNonLocalClass(ktDeclaration)
-            is KtNamedFunction -> findNonLocalFunction(ktDeclaration)
-            is KtProperty -> findNonLocalProperty(ktDeclaration)
-            is KtParameter -> findParameter(ktDeclaration)
-            is KtPropertyAccessor -> findNonLocalPropertyAccessor(ktDeclaration)
-            is KtTypeParameter -> findNonLocalTypeParameter(ktDeclaration)
+    fun findNonLocalDeclaration(ktDeclaration: KtDeclaration): FirDeclaration = when (ktDeclaration) {
+        is KtEnumEntry -> findNonLocalEnumEntry(ktDeclaration)
+        is KtClassLikeDeclaration -> findNonLocalClassLikeDeclaration(ktDeclaration)
+        is KtConstructor<*> -> findConstructorOfNonLocalClass(ktDeclaration)
+        is KtNamedFunction -> findNonLocalFunction(ktDeclaration)
+        is KtProperty -> findNonLocalProperty(ktDeclaration)
+        is KtParameter -> findParameter(ktDeclaration)
+        is KtPropertyAccessor -> findNonLocalPropertyAccessor(ktDeclaration)
+        is KtTypeParameter -> findNonLocalTypeParameter(ktDeclaration)
 
-            else -> errorWithFirSpecificEntries("Unsupported compiled declaration of type", psi = ktDeclaration)
+        else -> errorWithFirSpecificEntries("Unsupported compiled declaration of type", psi = ktDeclaration)
+    }
+
+    private fun FirSymbolProvider.findFunctionCandidates(function: KtNamedFunction): List<FirFunctionSymbol<*>> =
+        findCallableCandidates(function, function.isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()
+
+    private fun FirSymbolProvider.findPropertyCandidates(property: KtProperty): List<FirPropertySymbol> =
+        findCallableCandidates(property, property.isTopLevel).filterIsInstance<FirPropertySymbol>()
+
+    private fun FirSymbolProvider.findCallableCandidates(
+        declaration: KtCallableDeclaration,
+        isTopLevel: Boolean,
+    ): List<FirCallableSymbol<*>> {
+        val shortName = declaration.nameAsSafeName
+
+        if (isTopLevel) {
+            val packageFqName = declaration.containingKtFile.packageFqName
+
+            @OptIn(FirSymbolProviderInternals::class)
+            return when (this) {
+                is LLFirModuleWithDependenciesSymbolProvider -> buildList {
+                    getTopLevelDeserializedCallableSymbolsToWithoutDependencies(this, packageFqName, shortName, declaration)
+                    friendBuiltinsProvider?.getTopLevelCallableSymbolsTo(this, packageFqName, shortName)
+                }
+                else -> getTopLevelCallableSymbols(packageFqName, shortName)
+            }
+        }
+
+        val containingClass = declaration.containingClassOrObject?.let(::findNonLocalClassLikeDeclaration)
+            ?: errorWithFirSpecificEntries("No containing non-local declaration found for", psi = declaration)
+
+        val scope = session.declaredMemberScope(containingClass as FirClass, memberRequiredPhase = null)
+        return when (declaration) {
+            is KtProperty -> scope.getProperties(shortName)
+            is KtNamedFunction -> scope.getFunctions(shortName)
+            else -> errorWithFirSpecificEntries("Unexpected callable ${declaration::class.simpleName}") {
+                withEntry("isTopLevel", isTopLevel.toString())
+                withPsiEntry("declaration", declaration)
+            }
         }
     }
 
@@ -57,8 +98,9 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         val firTypeParameterRefOwner = firDeclaration as? FirTypeParameterRefsOwner ?: errorWithFirSpecificEntries(
             "No fir found by $owner",
             psi = owner,
-            fir = firDeclaration
+            fir = firDeclaration,
         )
+
         return firTypeParameterRefOwner.typeParameters.find { it.realPsi === param } as FirDeclaration
     }
 
@@ -75,21 +117,16 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
     }
 
     private fun findNonLocalEnumEntry(declaration: KtEnumEntry): FirEnumEntry {
-        require(!declaration.isLocal)
-        val classId = declaration.containingClassOrObject?.getClassId()
-            ?: errorWithFirSpecificEntries("Non-local class should have classId", psi = declaration)
+        val classCandidate = declaration.containingClassOrObject?.let(::findNonLocalClassLikeDeclaration)
+            ?: errorWithFirSpecificEntries("Enum entry must have containing class", psi = declaration)
 
-        val classCandidate = symbolProvider.getClassLikeSymbolByClassId(classId)
-            ?: errorWithFirSpecificEntries("We should be able to find a symbol for $classId", psi = declaration)
-
-        return (classCandidate.fir as? FirRegularClass)?.declarations?.first {
+        return (classCandidate as FirRegularClass).declarations.first {
             it is FirEnumEntry && it.realPsi === declaration
         } as FirEnumEntry
     }
 
     private fun findNonLocalClassLikeDeclaration(declaration: KtClassLikeDeclaration): FirClassLikeDeclaration {
-        val classId = declaration.getClassId()
-            ?: errorWithFirSpecificEntries("Non-local class should have classId", psi = declaration)
+        val classId = declaration.getClassId() ?: errorWithFirSpecificEntries("Non-local class should have classId", psi = declaration)
 
         val classCandidate = when (symbolProvider) {
             is LLFirModuleWithDependenciesSymbolProvider -> {
@@ -118,14 +155,9 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         val containingClass = declaration.containingClassOrObject
             ?: errorWithFirSpecificEntries("Constructor must have outer class", psi = declaration)
 
-        require(!containingClass.isLocal)
-        val classId = containingClass.getClassId()
-            ?: errorWithFirSpecificEntries("Non-local class should have classId", psi = declaration)
-
-        val constructorCandidate =
-            symbolProvider.getClassDeclaredConstructors(classId)
-                .singleOrNull { it.fir.realPsi === declaration }
-                ?: errorWithFirSpecificEntries("We should be able to find a constructor", psi = declaration)
+        val containingFirClass = findNonLocalClassLikeDeclaration(containingClass) as FirClass
+        val constructorCandidate = containingFirClass.constructors(symbolProvider.session).singleOrNull { it.fir.realPsi === declaration }
+            ?: errorWithFirSpecificEntries("We should be able to find a constructor", psi = declaration, fir = containingFirClass)
 
         return constructorCandidate.fir
     }
@@ -134,12 +166,10 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         require(!declaration.isLocal)
 
         val candidates = symbolProvider.findFunctionCandidates(declaration)
-        val functionCandidate =
-            candidates
-                .firstOrNull { it.fir.realPsi === declaration }
-                ?: errorWithFirSpecificEntries("We should be able to find a symbol for function", psi = declaration) {
-                    withCandidates(candidates)
-                }
+        val functionCandidate = candidates.firstOrNull { it.fir.realPsi === declaration }
+            ?: errorWithFirSpecificEntries("We should be able to find a symbol for function", psi = declaration) {
+                withCandidates(candidates)
+            }
 
         return functionCandidate.fir
     }
@@ -149,11 +179,10 @@ internal class FirDeclarationForCompiledElementSearcher(private val symbolProvid
         require(!declaration.isLocal)
 
         val candidates = symbolProvider.findPropertyCandidates(declaration)
-        val propertyCandidate =
-            candidates.firstOrNull { it.fir.realPsi === declaration }
-                ?: errorWithFirSpecificEntries("We should be able to find a symbol for property", psi = declaration) {
-                    withCandidates(candidates)
-                }
+        val propertyCandidate = candidates.firstOrNull { it.fir.realPsi === declaration }
+            ?: errorWithFirSpecificEntries("We should be able to find a symbol for property", psi = declaration) {
+                withCandidates(candidates)
+            }
 
         return propertyCandidate.fir
     }
@@ -181,38 +210,6 @@ private val LLFirModuleWithDependenciesSymbolProvider.friendBuiltinsProvider: Fi
         return null
     }
 
-private fun FirSymbolProvider.findFunctionCandidates(function: KtNamedFunction): List<FirFunctionSymbol<*>> =
-    findCallableCandidates(function, function.isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()
-
-private fun FirSymbolProvider.findPropertyCandidates(property: KtProperty): List<FirPropertySymbol> =
-    findCallableCandidates(property, property.isTopLevel).filterIsInstance<FirPropertySymbol>()
-
-private fun FirSymbolProvider.findCallableCandidates(
-    declaration: KtCallableDeclaration,
-    isTopLevel: Boolean
-): List<FirCallableSymbol<*>> {
-    val shortName = declaration.nameAsSafeName
-
-    if (isTopLevel) {
-        val packageFqName = declaration.containingKtFile.packageFqName
-
-        @OptIn(FirSymbolProviderInternals::class)
-        return when (this) {
-            is LLFirModuleWithDependenciesSymbolProvider -> buildList {
-                getTopLevelDeserializedCallableSymbolsToWithoutDependencies(this, packageFqName, shortName, declaration)
-                friendBuiltinsProvider?.getTopLevelCallableSymbolsTo(this, packageFqName, shortName)
-            }
-            else -> getTopLevelCallableSymbols(packageFqName, shortName)
-        }
-    }
-
-    val containerClassId = declaration.containingClassOrObject?.getClassId()
-        ?: errorWithFirSpecificEntries("No containing non-local declaration found for", psi = declaration)
-
-    return getClassDeclaredFunctionSymbols(containerClassId, shortName) +
-            getClassDeclaredPropertySymbols(containerClassId, shortName)
-}
-
 private fun ExceptionAttachmentBuilder.withCandidates(candidates: List<FirBasedSymbol<*>>) {
     withEntry("Candidates count", candidates.size.toString())
     for ((index, candidate) in candidates.withIndex()) {