From 0502da474b5037188e670846778d824f2a4b9c5b Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Thu, 11 Jan 2024 17:04:37 +0900
Subject: [PATCH] [Analysis API] Load scripting extensions for script dangling
 files

The commit fixes failing 'ContextCollectorScriptTestGenerated'
for dangling file copies.
---
 .../sessions/LLFirAbstractSessionFactory.kt   | 42 +++++++++++--------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirAbstractSessionFactory.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirAbstractSessionFactory.kt
index 809049441af..2e0680b70ac 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirAbstractSessionFactory.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/LLFirAbstractSessionFactory.kt
@@ -160,28 +160,32 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
             register(FirRegisteredPluginAnnotations::class, FirRegisteredPluginAnnotationsImpl(this))
             register(FirJvmTypeMapper::class, FirJvmTypeMapper(this))
 
-            FirSessionConfigurator(this).apply {
-                val hostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {}
-                val scriptDefinition = module.file.findScriptDefinition()
-                    ?: errorWithAttachment("Cannot load script definition") {
-                        withVirtualFileEntry("file", module.file.virtualFile)
-                    }
-
-                val extensionRegistrar = FirScriptingCompilerExtensionIdeRegistrar(
-                    project,
-                    hostConfiguration,
-                    scriptDefinitionSources = emptyList(),
-                    scriptDefinitions = listOf(scriptDefinition)
-                )
-
-                registerExtensions(extensionRegistrar.configure())
-                registerExtensions(FirScriptingSamWithReceiverExtensionRegistrar().configure())
-            }.configure()
+            registerScriptExtensions(this, module.file)
 
             LLFirSessionConfigurator.configure(this)
         }
     }
 
+    private fun registerScriptExtensions(session: LLFirSession, file: KtFile) {
+        FirSessionConfigurator(session).apply {
+            val hostConfiguration = ScriptingHostConfiguration(defaultJvmScriptingHostConfiguration) {}
+            val scriptDefinition = file.findScriptDefinition()
+                ?: errorWithAttachment("Cannot load script definition") {
+                    withVirtualFileEntry("file", file.virtualFile)
+                }
+
+            val extensionRegistrar = FirScriptingCompilerExtensionIdeRegistrar(
+                project,
+                hostConfiguration,
+                scriptDefinitionSources = emptyList(),
+                scriptDefinitions = listOf(scriptDefinition)
+            )
+
+            registerExtensions(extensionRegistrar.configure())
+            registerExtensions(FirScriptingSamWithReceiverExtensionRegistrar().configure())
+        }.configure()
+    }
+
     fun createNotUnderContentRootResolvableSession(module: KtNotUnderContentRootModule): LLFirNotUnderContentRootResolvableModuleSession {
         val builtinsSession = LLFirBuiltinsSessionFactory.getInstance(project).getBuiltinsSession(JvmPlatforms.unspecifiedJvmPlatform)
         val languageVersionSettings = ProjectStructureProvider.getInstance(project).globalLanguageVersionSettings
@@ -532,6 +536,10 @@ internal abstract class LLFirAbstractSessionFactory(protected val project: Proje
                 .createIfNeeded(this)
                 ?.also { register(FirSwitchableExtensionDeclarationsSymbolProvider::class, it) }
 
+            if (contextModule is KtScriptModule) {
+                registerScriptExtensions(this, contextModule.file)
+            }
+
             val context = DanglingFileSessionCreationContext(
                 moduleData,
                 firProvider,