Fix sequence of the script definitions search - explicit ones should be tried first

compiler/tests/org/jetbrains/kotlin/scripts/ScriptProviderTest.kt

@@ -22,7 +22,7 @@ class ScriptProviderTest : KtUsefulTestCase() {
         val genDefCounter = AtomicInteger()
         val standardDef = FakeScriptDefinition()
         val shadedDef = FakeScriptDefinition(".x.kts")
-        val provider = CliScriptDefinitionProvider().apply {
+        val provider = TestCliScriptDefinitionProvider(standardDef).apply {
             setScriptDefinitions(listOf(shadedDef, standardDef))
             setScriptDefinitionsSources(listOf(TestScriptDefinitionSource(genDefCounter, ".y.kts", ".x.kts")))
         }
@@ -41,8 +41,8 @@ class ScriptProviderTest : KtUsefulTestCase() {
 
         provider.isScript("a.x.kts").let {
             Assert.assertTrue(it)
-            Assert.assertEquals(2, genDefCounter.get())
+            Assert.assertEquals(1, genDefCounter.get())
-            Assert.assertEquals(0, shadedDef.matchCounter.get())
+            Assert.assertEquals(1, shadedDef.matchCounter.get())
         }
 
         provider.isScript("a.z.kts").let {
@@ -58,7 +58,7 @@ class ScriptProviderTest : KtUsefulTestCase() {
     }
 }
 
-private class FakeScriptDefinition(val suffix: String = ".kts") : KotlinScriptDefinition(ScriptTemplateWithArgs::class) {
+private open class FakeScriptDefinition(val suffix: String = ".kts") : KotlinScriptDefinition(ScriptTemplateWithArgs::class) {
     val matchCounter = AtomicInteger()
     override fun isScript(fileName: String): Boolean = fileName.endsWith(suffix).also {
         if (it) matchCounter.incrementAndGet()
@@ -76,4 +76,8 @@ private class TestScriptDefinitionSource(val counter: AtomicInteger, val defGens
             yield(gen())
         }
     }
+}
+
+private class TestCliScriptDefinitionProvider(private val standardDef: KotlinScriptDefinition) : CliScriptDefinitionProvider() {
+    override fun getDefaultScriptDefinition(): KotlinScriptDefinition = standardDef
 }

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/legacy/CliScriptDefinitionProvider.kt

@@ -10,12 +10,16 @@ import org.jetbrains.kotlin.script.ScriptDefinitionsSource
 import org.jetbrains.kotlin.script.StandardScriptDefinition
 import kotlin.concurrent.write
 
-class CliScriptDefinitionProvider : LazyScriptDefinitionProvider() {
+open class CliScriptDefinitionProvider : LazyScriptDefinitionProvider() {
     private val definitionsFromSources: MutableList<Sequence<KotlinScriptDefinition>> = arrayListOf()
-    private val definitions: MutableList<KotlinScriptDefinition> = arrayListOf(StandardScriptDefinition)
+    private val definitions: MutableList<KotlinScriptDefinition> = arrayListOf()
+    private var hasStandardDefinition = true
 
-    override val currentDefinitions: Sequence<KotlinScriptDefinition> =
+    override val currentDefinitions: Sequence<KotlinScriptDefinition>
-        definitionsFromSources.asSequence().flatMap { it } + definitions.asSequence()
+        get() {
+            val base = definitions.asSequence() + definitionsFromSources.asSequence().flatMap { it }
+            return if (hasStandardDefinition) base + getDefaultScriptDefinition() else base
+        }
 
     override fun getDefaultScriptDefinition(): KotlinScriptDefinition {
         return StandardScriptDefinition
@@ -24,7 +28,9 @@ class CliScriptDefinitionProvider : LazyScriptDefinitionProvider() {
     fun setScriptDefinitions(newDefinitions: List<KotlinScriptDefinition>) {
         lock.write {
             definitions.clear()
-            definitions.addAll(newDefinitions)
+            val (withoutStdDef, stdDef) = newDefinitions.partition { it != getDefaultScriptDefinition() }
+            definitions.addAll(withoutStdDef)
+            hasStandardDefinition = stdDef.isNotEmpty()
         }
     }