From bd6d492b0c437eaf942f21529667da77d205cbc5 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Sat, 23 Sep 2023 01:30:15 +0900
Subject: [PATCH] [LL API] Compile code fragment together with its context file

Recently, 'KtFirCompilerFacility' started to compile files from
dependent modules separately. Until recently, the code fragment was
in the same module as its context, so they naturally went in the same
compilation round. Now a separate 'KtCodeFragmentModule' is created,
so the code fragment case has to be handled explicitly.

This fixes the following groups of tests:
- FirIdeNormalAnalysisLibrarySourceModuleCompilerFacilityTestGenerated
- K2IdeK2CodeKotlinEvaluateExpressionTestGenerated (in IntelliJ)
---
 .../api/fir/components/KtFirCompilerFacility.kt     | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
index e1b4d17491b..b7d1f5f6a06 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCompilerFacility.kt
@@ -28,6 +28,8 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirMo
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
+import org.jetbrains.kotlin.analysis.project.structure.KtCodeFragmentModule
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
 import org.jetbrains.kotlin.backend.jvm.*
@@ -139,8 +141,8 @@ internal class KtFirCompilerFacility(
         }
 
         // Files in the code fragment context module are compiled together with the code fragment itself.
-        val targetModule = mainFirFile.llFirModuleData.ktModule
-        val (targetFiles, dependencyFiles) = filesToCompile.partition { firResolveSession.getModule(it) == targetModule }
+        val targetModules = computeTargetModules(mainFirFile.llFirModuleData.ktModule)
+        val (targetFiles, dependencyFiles) = filesToCompile.partition { firResolveSession.getModule(it) in targetModules }
         require(targetFiles.isNotEmpty())
 
         val jvmIrDeserializer = JvmIrDeserializerImpl()
@@ -244,6 +246,13 @@ internal class KtFirCompilerFacility(
         }
     }
 
+    private fun computeTargetModules(module: KtModule): List<KtModule> {
+        return when (module) {
+            is KtCodeFragmentModule -> listOf(module.contextModule, module)
+            else -> listOf(module)
+        }
+    }
+
     private fun runFir2Ir(
         session: LLFirSession,
         firFiles: List<FirFile>,