From b2b364eca68968f4f3d6316f5ed010fb0c947b91 Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Tue, 22 Dec 2020 14:37:20 +0100
Subject: [PATCH] FIR IDE: resolve fir file to BODY_RESOLVE when it requested
 by getOrBuildFir

---
 .../api/FirModuleResolveStateForCompletion.kt |  1 +
 .../level/api/FirModuleResolveStateImpl.kt    |  2 +-
 .../api/element/builder/FirElementBuilder.kt  | 29 ++++++++++++++++---
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
index 58f42126696..608cda97437 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateForCompletion.kt
@@ -47,6 +47,7 @@ internal class FirModuleResolveStateForCompletion(
         synchronized(completionMapping) { completionMapping[psi] }?.let { return it }
         return originalState.elementBuilder.getOrBuildFirFor(
             element,
+            originalState.firFileBuilder,
             originalState.rootModuleSession.cache,
             fileStructureCache,
         )
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt
index d84d7bbde1d..9691fd39509 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/FirModuleResolveStateImpl.kt
@@ -52,7 +52,7 @@ internal class FirModuleResolveStateImpl(
         sessionProvider.getSession(moduleInfo)!!
 
     override fun getOrBuildFirFor(element: KtElement): FirElement =
-        elementBuilder.getOrBuildFirFor(element, rootModuleSession.cache, fileStructureCache)
+        elementBuilder.getOrBuildFirFor(element, firFileBuilder, rootModuleSession.cache, fileStructureCache)
 
     override fun getFirFile(ktFile: KtFile): FirFile =
         firFileBuilder.buildRawFirFileWithCaching(ktFile, rootModuleSession.cache, lazyBodiesMode = false)
diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
index 4452f947608..506d43dc07f 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/element/builder/FirElementBuilder.kt
@@ -6,11 +6,13 @@
 package org.jetbrains.kotlin.idea.fir.low.level.api.element.builder
 
 import com.intellij.psi.PsiElement
-import com.intellij.psi.util.parentOfType
 import com.intellij.psi.util.parentsOfType
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.idea.fir.low.level.api.annotations.ThreadSafe
+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.file.structure.FileStructureCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.FileStructureElement
@@ -45,12 +47,31 @@ internal class FirElementBuilder {
 
     fun getOrBuildFirFor(
         element: KtElement,
+        firFileBuilder: FirFileBuilder,
         moduleFileCache: ModuleFileCache,
         fileStructureCache: FileStructureCache,
-    ): FirElement {
-        val fileStructure = fileStructureCache.getFileStructure(element.containingKtFile, moduleFileCache)
-        val mappings = fileStructure.getStructureElementFor(element).mappings
+    ): FirElement = when (element) {
+        is KtFile -> getOrBuildFirForKtFile(element, firFileBuilder, moduleFileCache)
+        else -> getOrBuildFirForNonKtFileElement(element, fileStructureCache, moduleFileCache)
+    }
 
+    private fun getOrBuildFirForKtFile(ktFile: KtFile, firFileBuilder: FirFileBuilder, moduleFileCache: ModuleFileCache): FirFile =
+        firFileBuilder.getFirFileResolvedToPhaseWithCaching(
+            ktFile,
+            moduleFileCache,
+            FirResolvePhase.BODY_RESOLVE,
+            checkPCE = true
+        )
+
+    private fun getOrBuildFirForNonKtFileElement(
+        element: KtElement,
+        fileStructureCache: FileStructureCache,
+        moduleFileCache: ModuleFileCache
+    ): FirElement {
+        require(element !is KtFile)
+        val fileStructure = fileStructureCache.getFileStructure(element.containingKtFile, moduleFileCache)
+
+        val mappings = fileStructure.getStructureElementFor(element).mappings
         val psi = getPsiAsFirElementSource(element)
         mappings[psi]?.let { return it }
         return psi.getFirOfClosestParent(mappings)?.second