FIR IDE: introduce withFirDeclaration function

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FirElementsRecorder
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -23,6 +24,7 @@ internal class FirModuleResolveStateForCompletion(
     private val originalState: FirModuleResolveStateImpl
 ) : FirModuleResolveState() {
     override val moduleInfo: IdeaModuleInfo get() = originalState.moduleInfo
+
     override val rootModuleSession get() = originalState.rootModuleSession
     override val firTransformerProvider: FirTransformerProvider get() = originalState.firTransformerProvider
     private val fileStructureCache = originalState.fileStructureCache
@@ -69,4 +71,8 @@ internal class FirModuleResolveStateForCompletion(
     override fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic> {
         error("Diagnostics should not be retrieved in completion")
     }
+
+    override fun findNonLocalSourceFirDeclaration(ktDeclaration: KtDeclaration): FirDeclaration {
+        error("Should not be used in completion")
+    }
 }

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.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState
 import org.jetbrains.kotlin.idea.fir.low.level.api.diagnostics.DiagnosticsCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirElementBuilder
@@ -20,8 +22,11 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataC
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+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.findSourceNonLocalFirDeclaration
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -56,6 +61,14 @@ internal class FirModuleResolveStateImpl(
         error("Should be called only from FirModuleResolveStateForCompletion")
     }
 
+    override fun findNonLocalSourceFirDeclaration(
+        ktDeclaration: KtDeclaration,
+    ): FirDeclaration = ktDeclaration.findSourceNonLocalFirDeclaration(
+        firFileBuilder,
+        rootModuleSession.firIdeProvider.symbolProvider,
+        sessionProvider.getModuleCache(ktDeclaration.getModuleInfo() as ModuleSourceInfo)
+    )
+
     override fun <D : FirDeclaration> resolvedFirToPhase(declaration: D, toPhase: FirResolvePhase): D {
         val fileCache = when (val session = declaration.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

@@ -12,11 +12,15 @@ 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.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 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.FirTransformerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
+import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 
@@ -37,6 +41,10 @@ abstract class FirModuleResolveState {
 
     internal abstract fun collectDiagnosticsForFile(ktFile: KtFile): Collection<Diagnostic>
 
+    abstract fun findNonLocalSourceFirDeclaration(
+        ktDeclaration: KtDeclaration,
+    ): FirDeclaration
+
     // todo temporary, used only in completion
     internal abstract fun recordPsiToFirMappingsForCompletionFrom(fir: FirDeclaration, firFile: FirFile, ktFile: KtFile)
 

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 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.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import kotlin.reflect.KClass
@@ -44,6 +45,25 @@ object LowLevelFirApiFacade {
     inline fun <R> withFirFile(ktFile: KtFile, resolveState: FirModuleResolveState, action: (FirFile) -> R): R =
         action(getFirFile(ktFile, resolveState))
 
+    /**
+     * Creates [FirDeclaration] by [KtDeclaration] and runs an [action] with it
+     * [FirDeclaration] passed to [action] should not be leaked outside [action] lambda
+     * [FirDeclaration] passed to [action] will be resolved at least to [phase]
+     * Otherwise, some threading problems may arise,
+     *
+     * [ktDeclaration] should be non-local declaration (should have fully qualified name)
+     */
+    inline fun <R> withFirDeclaration(
+        ktDeclaration: KtDeclaration,
+        resolveState: FirModuleResolveState,
+        phase: FirResolvePhase = FirResolvePhase.RAW_FIR,
+        action: (FirDeclaration) -> R
+    ): R {
+        val firDeclaration = resolveState.findNonLocalSourceFirDeclaration(ktDeclaration)
+        resolvedFirToPhase(firDeclaration, phase, resolveState)
+        return action(firDeclaration)
+    }
+
     fun getDiagnosticsFor(element: KtElement, resolveState: FirModuleResolveState): Collection<Diagnostic> =
         resolveState.getDiagnostics(element)
 

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.hasExplicitTypeOrUnit
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.replaceFirst
 import org.jetbrains.kotlin.idea.util.getElementTextInContext
@@ -111,7 +111,7 @@ internal class FileStructure(
     }
 
     private fun createDeclarationStructure(declaration: KtDeclaration): FileStructureElement {
-        val firDeclaration = declaration.findNonLocalFirDeclaration(firFileBuilder, firIdeProvider, moduleFileCache)
+        val firDeclaration = declaration.findSourceNonLocalFirDeclaration(firFileBuilder, firIdeProvider, moduleFileCache)
         firLazyDeclarationResolver.lazyResolveDeclaration(
             firDeclaration,
             moduleFileCache,

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/lazy/resolve/FirLazyDeclarationResolver.kt

@@ -9,7 +9,6 @@ import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
-import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirDesignatedBodyResolveTransformerForIDE
@@ -19,7 +18,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.checkCanceled
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
-import org.jetbrains.kotlin.idea.fir.low.level.api.util.findNonLocalFirDeclaration
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.getContainingFile
 import org.jetbrains.kotlin.psi.*
 
@@ -145,7 +144,7 @@ internal class FirLazyDeclarationResolver(
         if (!KtPsiUtil.isLocal(ktDeclaration)) return this
         val nonLocalPsi = ktDeclaration.getNonLocalContainingOrThisDeclaration()
             ?: error("Container for local declaration cannot be null")
-        return nonLocalPsi.findNonLocalFirDeclaration(firFileBuilder, provider, moduleFileCache)
+        return nonLocalPsi.findSourceNonLocalFirDeclaration(firFileBuilder, provider.symbolProvider, moduleFileCache)
     }
 }
 

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.caches.project.isLibraryClasses
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.Immutable
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 
 @Immutable
 class FirIdeSessionProvider internal constructor(
@@ -23,4 +24,7 @@ class FirIdeSessionProvider internal constructor(
 ) : FirSessionProvider {
     override fun getSession(moduleInfo: ModuleInfo): FirSession? =
         sessions[moduleInfo]
+
+    internal fun getModuleCache(moduleSourceInfo: ModuleSourceInfo): ModuleFileCache =
+        (sessions.getValue(moduleSourceInfo) as FirIdeSourcesSession).cache
 }

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

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.declarations.FirTypeAlias
 import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
+import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.util.classIdIfNonLocal
@@ -17,18 +17,18 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
-internal fun KtDeclaration.findNonLocalFirDeclaration(
+internal fun KtDeclaration.findSourceNonLocalFirDeclaration(
     firFileBuilder: FirFileBuilder,
-    provider: FirProvider,
+    firSymbolProvider: FirSymbolProvider,
     moduleFileCache: ModuleFileCache
 ): FirDeclaration {
     require(!KtPsiUtil.isLocal(this))
     return when {
-        this is KtClassOrObject -> findFir(provider)
+        this is KtClassOrObject -> findFir(firSymbolProvider)
         this is KtNamedDeclaration && (this is KtProperty || this is KtNamedFunction) -> {
             val containerClass = containingClassOrObject
             val declarations = if (containerClass != null) {
-                val containerClassFir = containerClass.findFir(provider)
+                val containerClassFir = containerClass.findFir(firSymbolProvider)
                 containerClassFir.declarations
             } else {
                 val ktFile = containingKtFile
@@ -39,11 +39,11 @@ internal fun KtDeclaration.findNonLocalFirDeclaration(
             declarations.first { it.psi === this || it.psi == original }
         }
         this is KtConstructor<*> -> {
-            val containerClassFir = containingClassOrObject?.findFir(provider)
+            val containerClassFir = containingClassOrObject?.findFir(firSymbolProvider)
                 ?: error("Container class should be not null for KtConstructor")
             containerClassFir.declarations.first { it.psi === this }
         }
-        this is KtTypeAlias -> findFir(provider)
+        this is KtTypeAlias -> findFir(firSymbolProvider)
         else -> error("Invalid container $this::class")
     }
 }
@@ -53,19 +53,19 @@ val ORIGINAL_DECLARATION_KEY = com.intellij.openapi.util.Key<KtDeclaration>("ORI
 var KtDeclaration.originalDeclaration by UserDataProperty(ORIGINAL_DECLARATION_KEY)
 
 
-private fun KtClassOrObject.findFir(provider: FirProvider): FirRegularClass {
+private fun KtClassOrObject.findFir(firSymbolProvider: FirSymbolProvider): FirRegularClass {
     val classId = classIdIfNonLocal()
         ?: error("Container classId should not be null for non-local declaration")
     return executeWithoutPCE {
-        provider.getFirClassifierByFqName(classId) as? FirRegularClass
+        firSymbolProvider.getClassLikeSymbolByFqName(classId)?.fir as? FirRegularClass
             ?: error("Could not find class $classId")
     }
 }
 
-private fun KtTypeAlias.findFir(provider: FirProvider): FirTypeAlias {
+private fun KtTypeAlias.findFir(firSymbolProvider: FirSymbolProvider): FirTypeAlias {
     val typeAlias = ClassId(containingKtFile.packageFqName, nameAsSafeName)
     return executeWithoutPCE {
-        provider.getFirClassifierByFqName(typeAlias) as? FirTypeAlias
+        firSymbolProvider.getClassLikeSymbolByFqName(typeAlias)?.fir as? FirTypeAlias
             ?: error("Could not find type alias $typeAlias")
     }
 }