From 6413294f3f076258905f3ed02a68a13f757d0b3f Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Tue, 1 Aug 2023 14:46:59 +0900
Subject: [PATCH] [Analysis API] Fix module provider for files from dependent
 sessions

Files coming from dependent sessions might be fully virtual, so there
won't be any backing 'VirtualFile'. However, they generally belong to
their original file's module.
---
 .../project/structure/KtStaticModuleProvider.kt    | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt
index 617b17ba20f..6c58a206769 100644
--- a/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt
+++ b/analysis/analysis-api-standalone/analysis-api-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/KtStaticModuleProvider.kt
@@ -18,16 +18,18 @@ class KtStaticModuleProvider(
 ) : KtStaticProjectStructureProvider() {
     @OptIn(KtModuleStructureInternals::class)
     override fun getModule(element: PsiElement, contextualModule: KtModule?): KtModule {
-        val containingFileAsPsiFile = element.containingFile
-        val containingFileAsVirtualFile = containingFileAsPsiFile.virtualFile
-        if (containingFileAsVirtualFile.extension == BuiltInSerializerProtocol.BUILTINS_FILE_EXTENSION) {
+        // Unwrap context-dependent file copies coming from dependent sessions
+        val containingPsiFile = element.containingFile.originalFile
+
+        val containingVirtualFile = containingPsiFile.virtualFile
+        if (containingVirtualFile.extension == BuiltInSerializerProtocol.BUILTINS_FILE_EXTENSION) {
             return builtinsModule
         }
 
-        containingFileAsVirtualFile.analysisExtensionFileContextModule?.let { return it }
+        containingVirtualFile.analysisExtensionFileContextModule?.let { return it }
 
         projectStructure.binaryModules.firstOrNull { binaryModule ->
-            containingFileAsVirtualFile in binaryModule.contentScope
+            containingVirtualFile in binaryModule.contentScope
         }?.let { return it }
 
         return projectStructure.mainModules.firstOrNull { module ->
@@ -35,7 +37,7 @@ class KtStaticModuleProvider(
         }?.ktModule
             ?: throw KotlinExceptionWithAttachments("Cannot find KtModule; see the attachment for more details.")
                 .withAttachment(
-                    containingFileAsVirtualFile.path,
+                    containingVirtualFile.path,
                     allKtModules.joinToString(separator = System.lineSeparator()) { it.asDebugString() }
                 )
     }