FIR IDE: Make findNonLocalDeclarationForCompiledElement a member of FirModuleResolveState

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api
 
 import com.intellij.openapi.project.Project
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
@@ -81,4 +80,8 @@ internal class FirModuleResolveStateDepended(
     @OptIn(InternalForInline::class)
     override fun findSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration =
         originalState.findSourceFirDeclaration(ktDeclaration)
+
+    @OptIn(InternalForInline::class)
+    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration) =
+        originalState.findSourceFirDeclaration(ktDeclaration)
 }

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

@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 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
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -29,11 +31,11 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarati
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSessionProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirIdeSourcesSession
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.FirElementFinder
+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.getParentOfType
+import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 internal class FirModuleResolveStateImpl(
     override val project: Project,
@@ -120,6 +122,66 @@ internal class FirModuleResolveStateImpl(
             ?: error("FirDeclaration was not found for\n${ktDeclaration.getElementTextInContext()}")
     }
 
+    @OptIn(InternalForInline::class)
+    override fun findSourceFirCompiledDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        require(ktDeclaration.containingKtFile.isCompiled) {
+            "This method will only work on compiled declarations, but this declaration is not compiled: ${ktDeclaration.getElementTextInContext()}"
+        }
+
+        return when (ktDeclaration) {
+            is KtClassOrObject -> {
+                require(!ktDeclaration.isLocal)
+                val classId = ktDeclaration.getClassId()
+                    ?: error("Non-local class should have classId. The class is ${ktDeclaration.getElementTextInContext()}")
+
+                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()}")
+        }
+    }
+
     override fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
         val fileCache = when (val session = declaration.moduleData.session) {
             is FirIdeSourcesSession -> session.cache

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

@@ -6,11 +6,9 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.api
 
 import com.intellij.openapi.project.Project
-import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -72,5 +70,16 @@ abstract class FirModuleResolveState {
         ktDeclaration: KtLambdaExpression,
     ): FirDeclaration
 
+    /**
+     * Looks for compiled non-local [ktDeclaration] declaration by querying its classId/callableId from the SymbolProvider.
+     *
+     * Works only if [ktDeclaration] is compiled (i.e. comes from .class file).
+     */
+    @InternalForInline
+    abstract fun findSourceFirCompiledDeclaration(
+        ktDeclaration: KtDeclaration
+    ): FirDeclaration
+
+
     internal abstract fun <D : FirDeclaration> resolveFirToPhase(declaration: D, toPhase: FirResolvePhase): D
 }

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

@@ -8,12 +8,10 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.api
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.resolve.symbolProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirIdeResolveStateService
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalDeclarationForCompiledElement
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
@@ -47,7 +45,7 @@ inline fun <R> KtDeclaration.withFirDeclaration(
     action: (FirDeclaration) -> R
 ): R {
     val firDeclaration = if (containingKtFile.isCompiled) {
-        findNonLocalDeclarationForCompiledElement(resolveState.rootModuleSession.symbolProvider)
+        resolveState.findSourceFirCompiledDeclaration(this)
     } else {
         resolveState.findSourceFirDeclaration(this)
     }

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

@@ -8,14 +8,12 @@ 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.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.fir.low.level.api.annotations.InternalForInline
 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
@@ -65,91 +63,24 @@ internal inline fun <reified F : FirDeclaration> KtDeclaration.findFirDeclaratio
     return fir
 }
 
-/**
+internal fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>): Boolean =
- * Looks for compiled non-local [this] declaration by querying its classId/callableId from the [symbolProvider].
- *
- * Works only if [this] is compiled (i.e. comes from .class file).
- */
-@InternalForInline
-fun KtDeclaration.findNonLocalDeclarationForCompiledElement(symbolProvider: FirSymbolProvider): FirDeclaration {
-    require(containingKtFile.isCompiled) {
-        "This method will only work on compiled declarations, but this declaration is not compiled: ${getElementTextInContext()}"
-    }
-
-    return when (this) {
-        is KtClassOrObject -> {
-            require(!isLocal)
-            val classId = getClassId()
-                ?: error("Non-local class should have classId. The class is ${getElementTextInContext()}")
-
-            val classCandidate = symbolProvider.getClassLikeSymbolByFqName(classId)
-                ?: error("We should be able to find a symbol for $classId")
-
-            classCandidate.fir
-        }
-
-        is KtConstructor<*> -> {
-            val containingClass = containingClassOrObject ?: error("Constructor must have outer class: ${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(this, it.fir) }
-                    ?: error("We should be able to find a constructor: ${getElementTextInContext()}")
-
-            constructorCandidate.fir
-        }
-
-        is KtNamedFunction -> {
-            require(!isLocal)
-
-            val functionCandidate =
-                findFunctionCandidates(symbolProvider, isTopLevel)
-                    .singleOrNull { representSameFunction(this, it.fir) }
-                    ?: error("We should be able to find a symbol for function ${name}: ${getElementTextInContext()}")
-
-            functionCandidate.fir
-        }
-
-        is KtProperty -> {
-            require(!isLocal)
-
-            val propertyCandidate =
-                findPropertyCandidates(symbolProvider, isTopLevel)
-                    .singleOrNull()
-                    ?: error("We should be able to find a symbol for property ${name}: ${getElementTextInContext()}")
-
-            propertyCandidate.fir
-        }
-
-        else -> error("Unsupported compiled declaration of type ${this::class}: ${getElementTextInContext()}")
-    }
-}
-
-private fun representSameFunction(psiFunction: KtNamedFunction, it: FirFunction<*>) =
     KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiFunction, it)
 
-private fun representSameConstructor(
+internal fun representSameConstructor(psiConstructor: KtConstructor<*>, firConstructor: FirConstructor): Boolean {
-    psiConstructor: KtConstructor<*>,
+    if ((firConstructor.isPrimary) != (psiConstructor is KtPrimaryConstructor)) {
-    firConstructor: FirConstructor
+        return false
-): Boolean {
-    if (firConstructor.isPrimary) {
-        return psiConstructor is KtPrimaryConstructor
     }
 
     return KtDeclarationAndFirDeclarationEqualityChecker.representsTheSameDeclaration(psiConstructor, firConstructor)
 }
 
-private fun KtCallableDeclaration.findFunctionCandidates(
+internal fun KtCallableDeclaration.findFunctionCandidates(
     symbolProvider: FirSymbolProvider,
     isTopLevel: Boolean
 ): List<FirFunctionSymbol<*>> =
     findCallableCandidates(symbolProvider, isTopLevel).filterIsInstance<FirFunctionSymbol<*>>()
 
-private fun KtCallableDeclaration.findPropertyCandidates(
+internal fun KtCallableDeclaration.findPropertyCandidates(
     symbolProvider: FirSymbolProvider,
     isTopLevel: Boolean
 ): List<FirPropertySymbol> =