From 6627b62c21d79fbfe5b1939e7821e5e7532a3674 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Tue, 18 Oct 2022 16:01:37 +0200
Subject: [PATCH] Fix sources extraction when processing extension is on

The latter is used e.g. for scripts in source roots handling
as implemented in the commit
"Scripting: update scripts in source roots handling"
The former implementation did not covered all the cases.

#KT-54355 fixed
---
 .../kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt |  6 +++++-
 .../cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt  |  7 +------
 .../compiler/plugin/ScriptingWithCliCompilerTest.kt  | 12 ++++++++++--
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt b/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt
index 99b5300d44f..cc911b080a7 100644
--- a/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt
+++ b/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt
@@ -418,7 +418,11 @@ class KotlinCoreEnvironment private constructor(
         return result
     }
 
-    fun getSourceFiles(): List<KtFile> = sourceFiles
+    fun getSourceFiles(): List<KtFile> =
+        ProcessSourcesBeforeCompilingExtension.getInstances(project)
+            .fold(sourceFiles as Collection<KtFile>) { files, extension ->
+                extension.processSources(files, configuration)
+            }.toList()
 
     internal fun report(severity: CompilerMessageSeverity, message: String) = configuration.report(severity, message)
 
diff --git a/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt b/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt
index e1905ccd1bb..1c3f9ceb833 100644
--- a/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt
+++ b/compiler/cli/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinToJVMBytecodeCompiler.kt
@@ -30,9 +30,7 @@ import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.config.CommonConfigurationKeys.LOOKUP_TRACKER
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporterFactory
 import org.jetbrains.kotlin.diagnostics.impl.BaseDiagnosticsCollector
-import org.jetbrains.kotlin.extensions.ProcessSourcesBeforeCompilingExtension
 import org.jetbrains.kotlin.ir.backend.jvm.jvmResolveLibraries
-import org.jetbrains.kotlin.load.java.JavaClassesTracker
 import org.jetbrains.kotlin.load.kotlin.ModuleVisibilityManager
 import org.jetbrains.kotlin.modules.Module
 import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
@@ -232,10 +230,7 @@ object KotlinToJVMBytecodeCompiler {
 
     fun analyze(environment: KotlinCoreEnvironment): AnalysisResult? {
         val collector = environment.messageCollector
-        val sourceFiles = ProcessSourcesBeforeCompilingExtension.getInstances(environment.project)
-            .fold(environment.getSourceFiles() as Collection<KtFile>) { files, extension ->
-                extension.processSources(files, environment.configuration)
-            }
+        val sourceFiles = environment.getSourceFiles()
 
         // Can be null for Scripts/REPL
         val performanceManager = environment.configuration.get(CLIConfigurationKeys.PERF_MANAGER)
diff --git a/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingWithCliCompilerTest.kt b/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingWithCliCompilerTest.kt
index d2a6be121c2..c83f38a787b 100644
--- a/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingWithCliCompilerTest.kt
+++ b/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingWithCliCompilerTest.kt
@@ -191,7 +191,7 @@ class ScriptingWithCliCompilerTest {
     @Test
     fun testCompileScriptWithRegularKotlin() {
 
-        fun compileVariant(vararg flags: String): Pair<List<String>, ExitCode> {
+        fun compileVariant(vararg flags: String, withScriptInstance: Boolean = true): Pair<List<String>, ExitCode> {
             return withTempDir { tmpdir ->
                 val (_, err, exitCode) = captureOutErrRet {
                     CLITool.doMainNoExit(
@@ -200,7 +200,10 @@ class ScriptingWithCliCompilerTest {
                             "-d", tmpdir.path,
                             "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
                             *flags,
-                            "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScriptInstance.kt",
+                            if (withScriptInstance)
+                                "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScriptInstance.kt"
+                            else
+                                "$TEST_DATA_DIR/compiler/mixedCompilation/nonScript.kt",
                             "$TEST_DATA_DIR/compiler/mixedCompilation/simpleScript.main.kts"
                         )
                     )
@@ -232,6 +235,11 @@ class ScriptingWithCliCompilerTest {
             Assert.assertEquals(ExitCode.COMPILATION_ERROR, exitCode)
         }
 
+        compileVariant("-language-version", "1.9", withScriptInstance = false).let { (errLines, exitCode) ->
+            Assert.assertTrue(errLines.none { it.startsWith(scriptInSourceRootWarning) })
+            Assert.assertEquals(ExitCode.OK, exitCode)
+        }
+
         compileVariant("-language-version", "1.9", "-Xallow-any-scripts-in-source-roots").let { (errLines, exitCode) ->
             Assert.assertTrue(errLines.none {
                 it.endsWith(unresolvedScriptError) || it.startsWith(scriptInSourceRootWarning)