FIR: Collect all scopes in FirModuleResolveStateImpl

We need to collect the scopes to later use them in reference shortening

Probably we will need to optimize this behaviour later, since it seems
inefficient

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

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.fir.FirSession
 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.psi
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.InternalForInline
@@ -112,4 +112,7 @@ internal class FirModuleResolveStateForCompletion(
     override fun getBuiltFirFileOrNull(ktFile: KtFile): FirFile? {
         error("Should not be used in completion")
     }
+
+    override fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext? =
+        originalState.getTowerDataContextForElement(element)
 }

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

@@ -12,8 +12,8 @@ import org.jetbrains.kotlin.fir.FirSession
 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.psi
 import org.jetbrains.kotlin.fir.realPsi
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
@@ -23,6 +23,7 @@ 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
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getClosestAvailableParentContext
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -44,7 +45,8 @@ internal class FirModuleResolveStateImpl(
     val firLazyDeclarationResolver: FirLazyDeclarationResolver,
 ) : FirModuleResolveState() {
     override val rootModuleSession: FirIdeSourcesSession get() = sessionProvider.rootModuleSession
-    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver)
+    private val collector = FirTowerDataContextCollector()
+    val fileStructureCache = FileStructureCache(firFileBuilder, firLazyDeclarationResolver, collector)
     val elementBuilder = FirElementBuilder()
     private val diagnosticsCollector = DiagnosticsCollector(fileStructureCache, rootModuleSession.cache)
 
@@ -103,7 +105,13 @@ internal class FirModuleResolveStateImpl(
         )
         if (container.resolvePhase < FirResolvePhase.BODY_RESOLVE) {
             val cache = (container.session as FirIdeSourcesSession).cache
-            firLazyDeclarationResolver.lazyResolveDeclaration(container, cache, FirResolvePhase.BODY_RESOLVE, checkPCE = false /*TODO*/)
+            firLazyDeclarationResolver.lazyResolveDeclaration(
+                container,
+                cache,
+                FirResolvePhase.BODY_RESOLVE,
+                checkPCE = false /*TODO*/,
+                towerDataContextCollector = collector,
+            )
         }
         val firDeclaration = FirElementFinder.findElementIn<FirDeclaration>(container) { firDeclaration ->
             when (val realPsi = firDeclaration.realPsi) {
@@ -126,7 +134,13 @@ internal class FirModuleResolveStateImpl(
             is FirIdeSourcesSession -> session.cache
             else -> return declaration
         }
-        firLazyDeclarationResolver.lazyResolveDeclaration(declaration, fileCache, toPhase, checkPCE = true)
+        firLazyDeclarationResolver.lazyResolveDeclaration(
+            declaration,
+            fileCache,
+            toPhase,
+            checkPCE = true,
+            towerDataContextCollector = collector,
+        )
         return declaration
     }
 
@@ -153,4 +167,8 @@ internal class FirModuleResolveStateImpl(
 
     override fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile? =
         cache.getContainerFirFile(declaration)
+
+    override fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext? {
+        return collector.getClosestAvailableParentContext(element)
+    }
 }

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

@@ -10,6 +10,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.resolve.FirTowerDataContext
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
 import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
@@ -78,4 +79,6 @@ abstract class FirModuleResolveState {
     )
 
     internal abstract fun getFirFile(declaration: FirDeclaration, cache: ModuleFileCache): FirFile?
+
+    abstract fun getTowerDataContextForElement(element: KtElement): FirTowerDataContext?
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirTowerDataContextCollector.kt

@@ -23,4 +23,12 @@ class FirTowerDataContextCollector {
     }
 
     fun getContext(psi: KtElement): FirTowerDataContext? = state[psi]
-}
+}
+
+fun FirTowerDataContextCollector.getClosestAvailableParentContext(element: KtElement): FirTowerDataContext? {
+    var current: KtElement = element
+    while (true) {
+        getContext(current)?.let { return it }
+        current = current.parent as? KtElement ?: return null
+    }
+}

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

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.FirTowerDataContextCollector
 import org.jetbrains.kotlin.idea.fir.low.level.api.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
@@ -24,7 +25,8 @@ internal class FileStructure(
     private val firFile: FirFile,
     private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
     private val firFileBuilder: FirFileBuilder,
-    private val moduleFileCache: ModuleFileCache
+    private val moduleFileCache: ModuleFileCache,
+    private val collector: FirTowerDataContextCollector
 ) {
     private val firIdeProvider = firFile.session.firIdeProvider
 
@@ -74,7 +76,8 @@ internal class FileStructure(
             firDeclaration,
             moduleFileCache,
             FirResolvePhase.BODY_RESOLVE,
-            checkPCE = true
+            checkPCE = true,
+            towerDataContextCollector = collector
         )
         return moduleFileCache.firFileLockProvider.withReadLock(firFile) {
             FileElementFactory.createFileStructureElement(firDeclaration, declaration, firFile)

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.fir.low.level.api.file.structure
 
+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.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
@@ -12,16 +13,17 @@ import org.jetbrains.kotlin.psi.KtFile
 import java.util.concurrent.ConcurrentHashMap
 
 /**
- * Belongs to a [org.jetbrains.kotlin.idea.fir.low.level.api.FirModuleResolveState]
+ * Belongs to a [org.jetbrains.kotlin.idea.fir.low.level.api.api.FirModuleResolveState]
  */
 internal class FileStructureCache(
     private val fileBuilder: FirFileBuilder,
     private val firLazyDeclarationResolver: FirLazyDeclarationResolver,
+    private val collector: FirTowerDataContextCollector,
 ) {
     private val cache = ConcurrentHashMap<KtFile, FileStructure>()
 
     fun getFileStructure(ktFile: KtFile, moduleFileCache: ModuleFileCache): FileStructure = cache.computeIfAbsent(ktFile) {
         val firFile = fileBuilder.buildRawFirFileWithCaching(ktFile, moduleFileCache, lazyBodiesMode = false)
-        FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache)
+        FileStructure(ktFile, firFile, firLazyDeclarationResolver, fileBuilder, moduleFileCache, collector)
     }
 }