From 432884fe84ba62ae47a9a346d4134df6b994b5f3 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Fri, 22 Dec 2023 01:09:59 +0900
Subject: [PATCH] [Analysis API] Use resolution components from the same
 session

In 'getFirForNonBodyElement()', a wrong 'FirProvider' was used for
elements from modules other than a use-site one, and it could lead to
"Modules are inconsistent" exceptions from 'LLFirFileBuilder'.
---
 .../fir/element/builder/FirElementBuilder.kt  | 31 +++++--------------
 .../state/LLFirResolvableResolveSession.kt    |  2 +-
 2 files changed, 9 insertions(+), 24 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
index 7635d27ab4e..d72e7e73b3e 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt
@@ -75,14 +75,11 @@ internal class FirElementBuilder(
         else -> true
     }
 
-    fun getOrBuildFirFor(
-        element: KtElement,
-        firResolveSession: LLFirResolveSession,
-    ): FirElement? {
+    fun getOrBuildFirFor(element: KtElement): FirElement? {
         return if (element is KtFile && element !is KtCodeFragment) {
             getOrBuildFirForKtFile(element)
         } else {
-            getFirForNonKtFileElement(element, firResolveSession)
+            getFirForNonKtFileElement(element)
         }
     }
 
@@ -92,18 +89,15 @@ internal class FirElementBuilder(
         return firFile
     }
 
-    private fun getFirForNonKtFileElement(
-        element: KtElement,
-        firResolveSession: LLFirResolveSession,
-    ): FirElement? {
+    private fun getFirForNonKtFileElement(element: KtElement): FirElement? {
         require(element !is KtFile || element is KtCodeFragment)
 
         if (!doKtElementHasCorrespondingFirElement(element)) {
             return null
         }
 
-        getFirForElementInsideAnnotations(element, firResolveSession)?.let { return it }
-        getFirForElementInsideTypes(element, firResolveSession)?.let { return it }
+        getFirForElementInsideAnnotations(element)?.let { return it }
+        getFirForElementInsideTypes(element)?.let { return it }
 
         val psi = getPsiAsFirElementSource(element) ?: return null
         val firFile = element.containingKtFile
@@ -116,7 +110,6 @@ internal class FirElementBuilder(
 
     private inline fun <T : KtElement> getFirForNonBodyElement(
         element: KtElement,
-        firResolveSession: LLFirResolveSession,
         anchorElementProvider: (KtElement) -> T?,
         declarationProvider: (T) -> KtDeclaration?,
         resolveAndFindFirForAnchor: (FirDeclaration, T) -> FirElement?,
@@ -128,7 +121,7 @@ internal class FirElementBuilder(
 
         val firDeclaration = nonLocalDeclaration.findSourceNonLocalFirDeclaration(
             firFileBuilder = moduleComponents.firFileBuilder,
-            provider = firResolveSession.useSiteFirSession.firProvider,
+            provider = moduleComponents.session.firProvider,
         )
 
         val anchorFir = resolveAndFindFirForAnchor(firDeclaration, anchorElement) ?: return null
@@ -144,23 +137,15 @@ internal class FirElementBuilder(
         return modifierList.owner as? KtDeclaration
     }
 
-    private fun getFirForElementInsideAnnotations(
-        element: KtElement,
-        firResolveSession: LLFirResolveSession,
-    ): FirElement? = getFirForNonBodyElement(
+    private fun getFirForElementInsideAnnotations(element: KtElement): FirElement? = getFirForNonBodyElement(
         element = element,
-        firResolveSession = firResolveSession,
         anchorElementProvider = { it.parentOfType<KtAnnotationEntry>(withSelf = true) },
         declarationProvider = { it.owner() },
         resolveAndFindFirForAnchor = { declaration, anchor -> declaration.resolveAndFindAnnotation(anchor, goDeep = true) },
     )
 
-    private fun getFirForElementInsideTypes(
-        element: KtElement,
-        firResolveSession: LLFirResolveSession,
-    ): FirElement? = getFirForNonBodyElement(
+    private fun getFirForElementInsideTypes(element: KtElement): FirElement? = getFirForNonBodyElement(
         element = element,
-        firResolveSession = firResolveSession,
         anchorElementProvider = { it.parentsOfType<KtTypeReference>(withSelf = true).lastOrNull() },
         declarationProvider = {
             when (val parent = it.parent) {
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt
index 1933669c03a..6d5c91f42ec 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/state/LLFirResolvableResolveSession.kt
@@ -45,7 +45,7 @@ internal class LLFirResolvableResolveSession(
 ) {
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
         val moduleComponents = getModuleComponentsForElement(element)
-        return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
+        return moduleComponents.elementsBuilder.getOrBuildFirFor(element)
     }
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile {