FIR IDE: Refactor findSourceFirCompiledDeclaration

Add `FirDeclarationForCompiledElementSearcher` class to encapsulate
searching by the compiled declarations, move it to the separate file

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt

@@ -9,11 +9,8 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.realPsi
-import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredConstructors
 import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
@@ -35,7 +32,6 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.util.*
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 internal class FirModuleResolveStateImpl(
     override val project: Project,
@@ -128,55 +124,13 @@ internal class FirModuleResolveStateImpl(
             "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
         }
 
+        val searcher = FirDeclarationForCompiledElementSearcher(rootModuleSession.symbolProvider)
+
         return when (ktDeclaration) {
-            is KtClassOrObject -> {
+            is KtClassOrObject -> searcher.findNonLocalClass(ktDeclaration)
-                require(!ktDeclaration.isLocal)
+            is KtConstructor<*> -> searcher.findConstructorOfNonLocalClass(ktDeclaration)
-                val classId = ktDeclaration.getClassId()
+            is KtNamedFunction -> searcher.findNonLocalFunction(ktDeclaration)
-                    ?: error("Non-local class should have classId. The class is ${ktDeclaration.getElementTextInContext()}")
+            is KtProperty -> searcher.findNonLocalProperty(ktDeclaration)
-
-                val classCandidate = rootModuleSession.symbolProvider.getClassLikeSymbolByFqName(classId)
-                    ?: error("We should be able to find a symbol for $classId")
-
-                classCandidate.fir
-            }
-
-            is KtConstructor<*> -> {
-                val containingClass = ktDeclaration.containingClassOrObject
-                    ?: error("Constructor must have outer class: ${ktDeclaration.getElementTextInContext()}")
-
-                require(!containingClass.isLocal)
-                val classId = containingClass.getClassId()
-                    ?: error("Non-local class should have classId. The class is ${containingClass.getElementTextInContext()}")
-
-                val constructorCandidate =
-                    rootModuleSession.symbolProvider.getClassDeclaredConstructors(classId)
-                        .singleOrNull { representSameConstructor(ktDeclaration, it.fir) }
-                        ?: error("We should be able to find a constructor: ${ktDeclaration.getElementTextInContext()}")
-
-                constructorCandidate.fir
-            }
-
-            is KtNamedFunction -> {
-                require(!ktDeclaration.isLocal)
-
-                val functionCandidate =
-                    ktDeclaration.findFunctionCandidates(rootModuleSession.symbolProvider, ktDeclaration.isTopLevel)
-                        .singleOrNull { representSameFunction(ktDeclaration, it.fir) }
-                        ?: error("We should be able to find a symbol for function ${ktDeclaration.name}: ${ktDeclaration.getElementTextInContext()}")
-
-                functionCandidate.fir
-            }
-
-            is KtProperty -> {
-                require(!ktDeclaration.isLocal)
-
-                val propertyCandidate =
-                    ktDeclaration.findPropertyCandidates(rootModuleSession.symbolProvider, ktDeclaration.isTopLevel)
-                        .singleOrNull()
-                        ?: error("We should be able to find a symbol for property ${ktDeclaration.name}: ${ktDeclaration.getElementTextInContext()}")
-
-                propertyCandidate.fir
-            }
 
             else -> error("Unsupported compiled declaration of type ${ktDeclaration::class}: ${ktDeclaration.getElementTextInContext()}")
         }

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

@@ -0,0 +1,109 @@
+/*
+ * Copyright 2010-2021 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.idea.fir.low.level.api.util
+
+import org.jetbrains.kotlin.fir.declarations.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredConstructors
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredFunctionSymbols
+import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredPropertySymbols
+import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.idea.fir.low.level.api.KtDeclarationAndFirDeclarationEqualityChecker
+import org.jetbrains.kotlin.idea.util.getElementTextInContext
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+
+/**
+ * Allows to search for FIR declarations by compiled [KtDeclaration]s.
+ */
+internal class FirDeclarationForCompiledElementSearcher(private val symbolProvider: FirSymbolProvider) {
+    fun findNonLocalClass(declaration: KtClassOrObject): FirClassLikeDeclaration<*> {
+        require(!declaration.isLocal)
+        val classId = declaration.getClassId()
+            ?: error("Non-local class should have classId. The class is ${declaration.getElementTextInContext()}")
+
+        val classCandidate = symbolProvider.getClassLikeSymbolByFqName(classId)
+            ?: error("We should be able to find a symbol for $classId")
+
+        return classCandidate.fir
+    }
+
+    fun findConstructorOfNonLocalClass(declaration: KtConstructor<*>): FirConstructor {
+        val containingClass = declaration.containingClassOrObject
+            ?: error("Constructor must have outer class: ${declaration.getElementTextInContext()}")
+
+        require(!containingClass.isLocal)
+        val classId = containingClass.getClassId()
+            ?: error("Non-local class should have classId. The class is ${containingClass.getElementTextInContext()}")
+
+        val constructorCandidate =
+            symbolProvider.getClassDeclaredConstructors(classId)
+                .singleOrNull { representSameConstructor(declaration, it.fir) }
+                ?: error("We should be able to find a constructor: ${declaration.getElementTextInContext()}")
+
+        return constructorCandidate.fir
+    }
+
+    fun findNonLocalFunction(declaration: KtNamedFunction): FirFunction<*> {
+        require(!declaration.isLocal)
+
+        val functionCandidate =
+            symbolProvider.findFunctionCandidates(declaration)
+                .singleOrNull { representSameFunction(declaration, it.fir) }
+                ?: error("We should be able to find a symbol for function ${declaration.name}: ${declaration.getElementTextInContext()}")
+
+        return functionCandidate.fir
+    }
+
+    fun findNonLocalProperty(declaration: KtProperty): FirProperty {
+        require(!declaration.isLocal)
+
+        val propertyCandidate =
+            symbolProvider.findPropertyCandidates(declaration)
+                .singleOrNull()
+                ?: error("We should be able to find a symbol for property ${declaration.name}: ${declaration.getElementTextInContext()}")
+
+        return propertyCandidate.fir
+    }
+
+}
+
+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<*>> {
+    if (isTopLevel) {
+        return getTopLevelCallableSymbols(declaration.containingKtFile.packageFqName, declaration.nameAsSafeName)
+    }
+
+    val containerClassId = declaration.containingClassOrObject?.getClassId()
+        ?: error("No containing non-local declaration found for ${declaration.getElementTextInContext()}")
+
+    return getClassDeclaredFunctionSymbols(containerClassId, declaration.nameAsSafeName) +
+            getClassDeclaredPropertySymbols(containerClassId, declaration.nameAsSafeName)
+}
+
+private fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>): Boolean =
+    KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiFunction, it)
+
+private fun representSameConstructor(psiConstructor: KtConstructor<*>, firConstructor: FirConstructor): Boolean {
+    if ((firConstructor.isPrimary) != (psiConstructor is KtPrimaryConstructor)) {
+        return false
+    }
+
+    return KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiConstructor, firConstructor)
+}

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/util/declarationUtils.kt

@@ -8,12 +8,6 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.util
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
-import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredFunctionSymbols
-import org.jetbrains.kotlin.fir.resolve.providers.getClassDeclaredPropertySymbols
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.idea.fir.low.level.api.KtDeclarationAndFirDeclarationEqualityChecker
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.InvalidFirElementTypeException
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
@@ -63,44 +57,6 @@ internal inline fun <reified F : FirDeclaration> KtDeclaration.findFirDeclaratio
     return fir
 }
 
-internal fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>): Boolean =
-    KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiFunction, it)
-
-internal fun representSameConstructor(psiConstructor: KtConstructor<*>, firConstructor: FirConstructor): Boolean {
-    if ((firConstructor.isPrimary) != (psiConstructor is KtPrimaryConstructor)) {
-        return false
-    }
-
-    return KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiConstructor, firConstructor)
-}
-
-internal fun KtCallableDeclaration.findFunctionCandidates(
-    symbolProvider: FirSymbolProvider,
-    isTopLevel: Boolean
-): List<FirFunctionSymbol<*>> =
-    findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()
-
-internal fun KtCallableDeclaration.findPropertyCandidates(
-    symbolProvider: FirSymbolProvider,
-    isTopLevel: Boolean
-): List<FirPropertySymbol> =
-    findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirPropertySymbol>()
-
-private fun KtCallableDeclaration.findCallableCandidates(
-    symbolProvider: FirSymbolProvider,
-    isTopLevel: Boolean
-): List<FirCallableSymbol<*>> {
-    if (isTopLevel) {
-        return symbolProvider.getTopLevelCallableSymbols(containingKtFile.packageFqName, nameAsSafeName)
-    }
-
-    val containerClassId = containingClassOrObject?.getClassId()
-        ?: error("No containing non-local declaration found for ${getElementTextInContext()}")
-
-    return symbolProvider.getClassDeclaredFunctionSymbols(containerClassId, nameAsSafeName) +
-            symbolProvider.getClassDeclaredPropertySymbols(containerClassId, nameAsSafeName)
-}
-
 private fun KtDeclaration.findSourceOfNonLocalFirDeclarationByTraversingWholeTree(
     firFileBuilder: FirFileBuilder,
     moduleFileCache: ModuleFileCache,