From 793cb02c7debc64dc3f05fb72319ca707dd81c2b Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Wed, 26 Apr 2023 18:06:05 +0200
Subject: [PATCH] [LL FIR] allow collecting tower data context for entire
 FirFile

Before, it was done only for import statements

^KT-57966
---
 .../api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt   | 9 ++++++++-
 .../fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt | 9 +++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
index 86ee42b2072..da707ffe781 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.analysis.low.level.api.fir.DeclarationCopyBuilder.withBodyFrom
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionDepended
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirWholeFileResolveTarget
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerDataContextAllElementsCollector
@@ -116,7 +117,13 @@ object LowLevelFirApiFacadeForResolveOnAir {
         require(firResolveSession is LLFirResolvableResolveSession)
 
         return if (place is KtFile) {
-            FileTowerProvider(place, onAirGetTowerContextForFile(firResolveSession, place))
+            FirTowerDataContextAllElementsCollector().also { collector ->
+                val session = firResolveSession.getSessionFor(place.getKtModule(firResolveSession.project)) as LLFirResolvableModuleSession
+                val moduleComponents = session.moduleComponents
+                val onAirFirFile = RawFirNonLocalDeclarationBuilder.buildNewFile(session, session.kotlinScopeProvider, place)
+                val target = LLFirWholeFileResolveTarget(onAirFirFile)
+                moduleComponents.firModuleLazyDeclarationResolver.lazyResolveTarget(target, FirResolvePhase.BODY_RESOLVE, collector)
+            }
         } else {
             val validPlace = PsiTreeUtil.findFirstParent(place, false) {
                 RawFirReplacement.isApplicableForReplacement(it as KtElement)
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt
index f2797dad868..99ab207b446 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/RawFirNonLocalDeclarationBuilder.kt
@@ -61,6 +61,15 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
     }
 
     companion object {
+        fun buildNewFile(
+            session: FirSession,
+            scopeProvider: FirScopeProvider,
+            file: KtFile,
+        ): FirFile {
+            val builder = RawFirBuilder(session,scopeProvider, bodyBuildingMode = BodyBuildingMode.LAZY_BODIES)
+            return builder.buildFirFile(file)
+        }
+
         fun buildNewSimpleFunction(
             session: FirSession,
             scopeProvider: FirScopeProvider,