K2 scripting: separate scripts compilation into another session

when scripts are compiled along with other sources.
#KT-65865 fixed

compiler/cli/src/org/jetbrains/kotlin/cli/common/FirSessionConstructionUtils.kt

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.cli.common
 import org.jetbrains.kotlin.KtSourceFile
 import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
 import org.jetbrains.kotlin.config.*
-import org.jetbrains.kotlin.fir.DependencyListForCliModule
+import org.jetbrains.kotlin.fir.*
-import org.jetbrains.kotlin.fir.FirModuleData
-import org.jetbrains.kotlin.fir.FirModuleDataImpl
-import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.OptInLanguageVersionSettingsCheckers
 import org.jetbrains.kotlin.fir.checkers.registerExtendedCommonCheckers
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrar
@@ -68,6 +65,7 @@ fun <F> prepareJvmSessions(
     librariesScope: AbstractProjectFileSearchScope,
     libraryList: DependencyListForCliModule,
     isCommonSource: (F) -> Boolean,
+    isScript: (F) -> Boolean,
     fileBelongsToModule: (F, String) -> Boolean,
     createProviderAndScopeForIncrementalCompilation: (List<F>) -> IncrementalCompilationContext?,
 ): List<SessionWithSources<F>> {
@@ -79,7 +77,7 @@ fun <F> prepareJvmSessions(
 
     return prepareSessions(
         files, configuration, rootModuleName, JvmPlatforms.unspecifiedJvmPlatform,
-        JvmPlatformAnalyzerServices, metadataCompilationMode = false, libraryList, isCommonSource, fileBelongsToModule,
+        JvmPlatformAnalyzerServices, metadataCompilationMode = false, libraryList, isCommonSource, isScript, fileBelongsToModule,
         createLibrarySession = { sessionProvider ->
             FirJvmSessionFactory.createLibrarySession(
                 rootModuleName,
@@ -149,7 +147,8 @@ fun <F> prepareJsSessions(
 ): List<SessionWithSources<F>> {
     return prepareSessions(
         files, configuration, rootModuleName, JsPlatforms.defaultJsPlatform, JsPlatformAnalyzerServices,
-        metadataCompilationMode = false, libraryList, isCommonSource, fileBelongsToModule,
+        metadataCompilationMode = false, libraryList, isCommonSource, isScript = { false },
+        fileBelongsToModule,
         createLibrarySession = { sessionProvider ->
             FirJsSessionFactory.createLibrarySession(
                 rootModuleName,
@@ -196,7 +195,8 @@ fun <F> prepareNativeSessions(
 ): List<SessionWithSources<F>> {
     return prepareSessions(
         files, configuration, rootModuleName, NativePlatforms.unspecifiedNativePlatform, NativePlatformAnalyzerServices,
-        metadataCompilationMode, libraryList, isCommonSource, fileBelongsToModule, createLibrarySession = { sessionProvider ->
+        metadataCompilationMode, libraryList, isCommonSource, isScript = { false },
+        fileBelongsToModule, createLibrarySession = { sessionProvider ->
             FirNativeSessionFactory.createLibrarySession(
                 rootModuleName,
                 resolvedLibraries,
@@ -244,7 +244,8 @@ fun <F> prepareWasmSessions(
     }
     return prepareSessions(
         files, configuration, rootModuleName, WasmPlatforms.Default, analyzerServices,
-        metadataCompilationMode = false, libraryList, isCommonSource, fileBelongsToModule,
+        metadataCompilationMode = false, libraryList, isCommonSource, isScript = { false },
+        fileBelongsToModule,
         createLibrarySession = { sessionProvider ->
             FirWasmSessionFactory.createLibrarySession(
                 rootModuleName,
@@ -291,7 +292,8 @@ fun <F> prepareCommonSessions(
 ): List<SessionWithSources<F>> {
     return prepareSessions(
         files, configuration, rootModuleName, CommonPlatforms.defaultCommonPlatform, CommonPlatformAnalyzerServices,
-        metadataCompilationMode = true, libraryList, isCommonSource, fileBelongsToModule, createLibrarySession = { sessionProvider ->
+        metadataCompilationMode = true, libraryList, isCommonSource, isScript = { false }, fileBelongsToModule,
+        createLibrarySession = { sessionProvider ->
             FirCommonSessionFactory.createLibrarySession(
                 rootModuleName,
                 sessionProvider,
@@ -335,11 +337,13 @@ private inline fun <F> prepareSessions(
     metadataCompilationMode: Boolean,
     libraryList: DependencyListForCliModule,
     isCommonSource: (F) -> Boolean,
+    isScript: (F) -> Boolean,
     fileBelongsToModule: (F, String) -> Boolean,
     createLibrarySession: (FirProjectSessionProvider) -> FirSession,
     createSourceSession: FirSessionProducer<F>,
 ): List<SessionWithSources<F>> {
     val languageVersionSettings = configuration.languageVersionSettings
+    val (scripts, nonScriptFiles) = files.partition(isScript)
 
     val isMppEnabled = languageVersionSettings.supportsFeature(LanguageFeature.MultiPlatformProjects)
     val hmppModuleStructure = configuration.get(CommonConfigurationKeys.HMPP_MODULE_STRUCTURE)
@@ -353,26 +357,58 @@ private inline fun <F> prepareSessions(
         }
     }
 
-    return when {
+    val nonScriptSessions =when {
-        metadataCompilationMode || !isMppEnabled -> listOf(
+        metadataCompilationMode || !isMppEnabled -> {
-            createSingleSession(
+            listOf(
-                files, rootModuleName, libraryList, targetPlatform, analyzerServices,
+                createSingleSession(
-                sessionProvider, sessionConfigurator, createSourceSession
+                    nonScriptFiles, rootModuleName, libraryList, targetPlatform, analyzerServices,
+                    sessionProvider, sessionConfigurator, createSourceSession
+                )
             )
-        )
+        }
 
         hmppModuleStructure == null -> createSessionsForLegacyMppProject(
-            files, rootModuleName, libraryList, targetPlatform, analyzerServices,
+            nonScriptFiles, rootModuleName, libraryList, targetPlatform, analyzerServices,
             sessionProvider, sessionConfigurator, isCommonSource, createSourceSession
         )
 
         else -> createSessionsForHmppProject(
-            files, rootModuleName, hmppModuleStructure, libraryList, targetPlatform, analyzerServices,
+            nonScriptFiles, rootModuleName, hmppModuleStructure, libraryList, targetPlatform, analyzerServices,
             sessionProvider, sessionConfigurator, fileBelongsToModule, createSourceSession
         )
     }
+
+    return if (scripts.isEmpty()) nonScriptSessions
+    else nonScriptSessions +
+            createScriptsSession(
+                scripts, rootModuleName, libraryList, nonScriptSessions.last().session.moduleData,
+                targetPlatform, analyzerServices, sessionProvider, sessionConfigurator, createSourceSession
+            )
 }
 
+private inline fun <F> createScriptsSession(
+    scripts: List<F>,
+    rootModuleName: Name,
+    libraryList: DependencyListForCliModule,
+    lastModuleData: FirModuleData,
+    targetPlatform: TargetPlatform,
+    analyzerServices: PlatformDependentAnalyzerServices,
+    sessionProvider: FirProjectSessionProvider,
+    noinline sessionConfigurator: FirSessionConfigurator.() -> Unit,
+    createSourceSession: FirSessionProducer<F>,
+): SessionWithSources<F> =
+    createSingleSession(
+        scripts, Name.identifier("${rootModuleName.asString()}-scripts"),
+        DependencyListForCliModule(
+            libraryList.regularDependencies,
+            listOf(lastModuleData),
+            libraryList.friendsDependencies,
+            libraryList.moduleDataProvider
+        ),
+        targetPlatform, analyzerServices,
+        sessionProvider, sessionConfigurator, createSourceSession
+    )
+
 private inline fun <F> createSingleSession(
     files: List<F>,
     rootModuleName: Name,

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/FirKotlinToJvmBytecodeCompiler.kt

@@ -214,6 +214,7 @@ object FirKotlinToJvmBytecodeCompiler {
             ktFiles, configuration, projectEnvironment, Name.special("<$rootModuleName>"),
             extensionRegistrars, librariesScope, libraryList,
             isCommonSource = { it.isCommonSource == true },
+            isScript = { it.isScript() },
             fileBelongsToModule = { file, moduleName -> file.hmppModuleName == moduleName },
             createProviderAndScopeForIncrementalCompilation = { providerAndScopeForIncrementalCompilation }
         )

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/pipeline/compilerPipeline.kt

@@ -296,6 +296,7 @@ fun compileModuleToAnalyzedFir(
         allSources, moduleConfiguration, projectEnvironment, Name.special("<$rootModuleName>"),
         extensionRegistrars, librariesScope, libraryList,
         isCommonSource = input.groupedSources.isCommonSourceForLt,
+        isScript = { false },
         fileBelongsToModule = input.groupedSources.fileBelongsToModuleForLt,
         createProviderAndScopeForIncrementalCompilation = { files ->
             val scope = projectEnvironment.getSearchScopeBySourceFiles(files)

compiler/fir/analysis-tests/test/org/jetbrains/kotlin/test/LoadedMetadataDumpHandler.kt

@@ -82,6 +82,7 @@ class JvmLoadedMetadataDumpHandler(testServices: TestServices) : AbstractLoadedM
             environment.getSearchScopeForProjectLibraries(),
             libraryList,
             isCommonSource = { false },
+            isScript = { false },
             fileBelongsToModule = { _, _ -> false },
             createProviderAndScopeForIncrementalCompilation = { null }
         )

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/ScriptJvmCompilerImpls.kt

@@ -365,7 +365,8 @@ private fun doCompileWithK2(
     )
     val session = prepareJvmSessions(
         sourceFiles, kotlinCompilerConfiguration, projectEnvironment, Name.special("<$rootModuleName>"), extensionRegistrars,
-        librariesScope, libraryList, isCommonSourceForPsi, fileBelongsToModuleForPsi,
+        librariesScope, libraryList, isCommonSourceForPsi, { false },
+        fileBelongsToModuleForPsi,
         createProviderAndScopeForIncrementalCompilation = { files ->
             createContextForIncrementalCompilation(
                 compilerInput.configuration,

plugins/scripting/scripting-compiler/testData/compiler/mixedCompilation/nonScript.kt

@@ -1,4 +1,6 @@
 
+package foo
+
 fun main() {
     println("OK")
 }

plugins/scripting/scripting-compiler/testData/compiler/mixedCompilation/nonScriptAccessingScript.kt

@@ -0,0 +1,5 @@
+
+fun main() {
+    println(SimpleScript_main().ok)
+    println(ok)
+}

plugins/scripting/scripting-compiler/testData/compiler/mixedCompilation/scriptAccessingNonScript.main.kts

@@ -0,0 +1,2 @@
+
+foo.main()

plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingWithCliCompilerTest.kt

@@ -13,12 +13,14 @@ import org.jetbrains.kotlin.scripting.compiler.test.linesSplitTrim
 import org.junit.Assert
 import org.junit.Test
 import java.io.File
+import java.net.URLClassLoader
 import java.nio.file.Files
 
 class ScriptingWithCliCompilerTest {
 
     companion object {
         const val TEST_DATA_DIR = "plugins/scripting/scripting-compiler/testData"
+        val SIMPLE_TEST_SCRIPT = "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScript.main.kts"
     }
 
     init {
@@ -196,7 +198,7 @@ class ScriptingWithCliCompilerTest {
                                 "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScriptInstance.kt"
                             else
                                 "$TEST_DATA_DIR/compiler/mixedCompilation/nonScript.kt",
-                            "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScript.main.kts"
+                            SIMPLE_TEST_SCRIPT
                         )
                     )
                 }
@@ -240,6 +242,54 @@ class ScriptingWithCliCompilerTest {
         }
     }
 
+    @Test
+    fun testAccessRegularSourceFromScript() {
+        withTempDir { tmpdir ->
+            val scriptPath = "$TEST_DATA_DIR/compiler/mixedCompilation/scriptAccessingNonScript.main.kts"
+            val ret =
+                CLITool.doMainNoExit(
+                    K2JVMCompiler(),
+                    arrayOf(
+                        "-P", "plugin:kotlin.scripting:disable-script-definitions-autoloading=true",
+                        "-cp", getMainKtsClassPath().joinToString(File.pathSeparator), "-d", tmpdir.path,
+                        "-Xuse-fir-lt=false",
+                        "-Xallow-any-scripts-in-source-roots", "-verbose",
+                        "$TEST_DATA_DIR/compiler/mixedCompilation/nonScript.kt",
+                        scriptPath,
+                    )
+                )
+            Assert.assertEquals(ExitCode.OK.code, ret.code)
+            val (out, _, _) = captureOutErrRet {
+                val cl = URLClassLoader((getMainKtsClassPath() + tmpdir).map { it.toURI().toURL() }.toTypedArray())
+                val klass = cl.loadClass("ScriptAccessingNonScript_main")
+                val ctor = klass.constructors.single()
+                ctor.newInstance(arrayOf<String>(), File(scriptPath))
+            }
+            Assert.assertEquals("OK", out.trim())
+        }
+    }
+
+    @Test
+    fun testAccessScriptFromRegularSource() {
+        withTempDir { tmpdir ->
+            val (_, err, ret) = captureOutErrRet {
+                CLITool.doMainNoExit(
+                    K2JVMCompiler(),
+                    arrayOf(
+                        "-P", "plugin:kotlin.scripting:disable-script-definitions-autoloading=true",
+                        "-cp", getMainKtsClassPath().joinToString(File.pathSeparator), "-d", tmpdir.path,
+                        "-Xuse-fir-lt=false",
+                        "-Xallow-any-scripts-in-source-roots", "-verbose",
+                        "$TEST_DATA_DIR/compiler/mixedCompilation/nonScriptAccessingScript.kt",
+                        SIMPLE_TEST_SCRIPT
+                    )
+                )
+            }
+            println(err)
+            Assert.assertEquals(ExitCode.COMPILATION_ERROR.code, ret.code)
+        }
+    }
+
     @Test
     fun testScriptMainKtsDiscovery() {
         withTempDir { tmpdir ->
@@ -256,7 +306,7 @@ class ScriptingWithCliCompilerTest {
                         )
                     )
                 }
-                Assert.assertEquals(0, ret.code)
+                Assert.assertEquals(ExitCode.OK.code, ret.code)
                 return err.linesSplitTrim()
             }
 
@@ -265,7 +315,7 @@ class ScriptingWithCliCompilerTest {
             val res1 = compileSuccessfullyGetStdErr("$TEST_DATA_DIR/compiler/mixedCompilation/nonScript.kt")
             Assert.assertTrue(res1.none { it.startsWith(loadMainKtsMessage) })
 
-            val res2 = compileSuccessfullyGetStdErr("$TEST_DATA_DIR/compiler/mixedCompilation/simpleScript.main.kts")
+            val res2 = compileSuccessfullyGetStdErr(SIMPLE_TEST_SCRIPT)
             Assert.assertTrue(res2.any { it.startsWith(loadMainKtsMessage) })
         }
     }