From 9a7d1948a7a6e8aacf922cf0c2dd7a2e5a5020bd Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Wed, 25 Nov 2020 21:28:07 +0100
Subject: [PATCH] Implement support for -Xdefault-script-extension cli option

---
 .../arguments/K2JVMCompilerArguments.kt       |  7 ++
 compiler/testData/cli/jvm/extraHelp.out       |  2 +
 .../cli/jvm/wrongScriptWithKtSource.out       |  2 +-
 .../scripting/definitions/ScriptDefinition.kt |  6 +-
 .../AbstractScriptEvaluationExtension.kt      | 48 ++++++++---
 .../compiler/plugin/impl/errorReporting.kt    |  3 +-
 .../plugin/ScriptingWithCliCompilerTest.kt    | 80 +++++++++++++++++++
 .../scripting/compiler/plugin/testUtil.kt     | 22 +++--
 8 files changed, 148 insertions(+), 22 deletions(-)

diff --git a/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt b/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt
index 1ce3fb3efa4..0c4252ef04c 100644
--- a/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt
+++ b/compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt
@@ -308,6 +308,13 @@ class K2JVMCompilerArguments : CommonCompilerArguments() {
     )
     var jvmDefault: String by FreezableVar(JvmDefaultMode.DEFAULT.description)
 
+    @Argument(
+        value = "-Xdefault-script-extension",
+        valueDescription = "<script filename extension>",
+        description = "Compile expressions and unrecognized scripts passed with the -script argument as scripts with given filename extension"
+    )
+    var defaultScriptExtension: String? by FreezableVar(null)
+
     @Argument(value = "-Xdisable-standard-script", description = "Disable standard kotlin script support")
     var disableStandardScript: Boolean by FreezableVar(false)
 
diff --git a/compiler/testData/cli/jvm/extraHelp.out b/compiler/testData/cli/jvm/extraHelp.out
index 0146bea18d5..836cb96f001 100644
--- a/compiler/testData/cli/jvm/extraHelp.out
+++ b/compiler/testData/cli/jvm/extraHelp.out
@@ -23,6 +23,8 @@ where advanced options include:
                              default is 'disable' in language version 1.2 and below,
                              'enable' since language version 1.3
   -Xdump-declarations-to=<path> Path to JSON file to dump Java to Kotlin declaration mappings
+  -Xdefault-script-extension=<script filename extension>
+                             Compile expressions and unrecognized scripts passed with the -script argument as scripts with given filename extension
   -Xdisable-standard-script  Disable standard kotlin script support
   -Xir-do-not-clear-binding-context
                              When using the IR backend, do not clear BindingContext between psi2ir and lowerings
diff --git a/compiler/testData/cli/jvm/wrongScriptWithKtSource.out b/compiler/testData/cli/jvm/wrongScriptWithKtSource.out
index 174f8829fbf..b0f33b4560a 100644
--- a/compiler/testData/cli/jvm/wrongScriptWithKtSource.out
+++ b/compiler/testData/cli/jvm/wrongScriptWithKtSource.out
@@ -1,2 +1,2 @@
-error: unrecognized script file: $TESTDATA_DIR$/simple.kt; Specify path to the script file as the first argument
+error: unrecognized script type: simple.kt; Specify path to the script file as the first argument
 COMPILATION_ERROR
diff --git a/plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/definitions/ScriptDefinition.kt b/plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/definitions/ScriptDefinition.kt
index 8892165da76..393d37642b2 100644
--- a/plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/definitions/ScriptDefinition.kt
+++ b/plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/definitions/ScriptDefinition.kt
@@ -11,7 +11,8 @@ import org.jetbrains.kotlin.scripting.resolve.KotlinScriptDefinitionFromAnnotate
 import java.io.File
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
-import kotlin.script.experimental.host.*
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.host.createScriptDefinitionFromTemplate
 import kotlin.script.experimental.jvm.baseClassLoader
 import kotlin.script.experimental.jvm.jvm
 
@@ -139,8 +140,9 @@ abstract class ScriptDefinition : UserDataHolderBase() {
         }
 
         override fun isScript(script: SourceCode): Boolean {
+            val extension = ".$fileExtension"
             val location = script.locationId ?: return false
-            return location.endsWith(".$fileExtension") && filePathPattern?.let {
+            return (script.name?.endsWith(extension) == true || location.endsWith(extension)) && filePathPattern?.let {
                 Regex(it).matches(FileUtilRt.toSystemIndependentName(location))
             } != false
         }
diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/AbstractScriptEvaluationExtension.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/AbstractScriptEvaluationExtension.kt
index c21b0b850dd..8d34a560680 100644
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/AbstractScriptEvaluationExtension.kt
+++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/AbstractScriptEvaluationExtension.kt
@@ -19,8 +19,10 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.scripting.configuration.ScriptingConfigurationKeys
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider
 import java.io.File
+import java.io.Serializable
 import java.util.*
 import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileScriptSource
 import kotlin.script.experimental.host.StringScriptSource
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvm.util.renderError
@@ -53,23 +55,21 @@ abstract class AbstractScriptEvaluationExtension : ScriptEvaluationExtension {
 
         setupScriptConfiguration(configuration)
 
+        val defaultScriptExtension =
+            (arguments as? K2JVMCompilerArguments)?.defaultScriptExtension?.let { if (it.startsWith('.')) it else ".$it" }
+
         val script = when {
             arguments is K2JVMCompilerArguments && arguments.expression != null -> {
-                StringScriptSource(arguments.expression!!, "script.kts")
+                StringScriptSource(arguments.expression!!, "script${defaultScriptExtension ?: ".kts"}")
             }
             arguments.script -> {
-                val scriptFile = File(arguments.freeArgs.first())
-                val script = scriptFile.toScriptSource()
+                val scriptFile = File(arguments.freeArgs.first()).normalize()
 
-                val error = when {
-                    !scriptFile.exists() -> "Script file not found: ${arguments.freeArgs.first()}"
-                    scriptFile.isDirectory  -> "Script argument points to a directory: ${arguments.freeArgs.first()}"
-                    !scriptDefinitionProvider.isScript(script) -> "Unrecognized script file: ${arguments.freeArgs.first()}"
-                    else -> null
-                }
-                if (error != null) {
+                fun invalidScript(error: String): ExitCode {
                     val extensionHint =
-                        if (configuration.get(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS)?.let { it.size == 1 && it.first().isDefault } == true) " (.kts)"
+                        if (configuration.get(ScriptingConfigurationKeys.SCRIPT_DEFINITIONS)
+                                ?.let { it.size == 1 && it.first().isDefault } == true
+                        ) " (.kts)"
                         else ""
                     messageCollector.report(
                         CompilerMessageSeverity.ERROR,
@@ -77,7 +77,22 @@ abstract class AbstractScriptEvaluationExtension : ScriptEvaluationExtension {
                     )
                     return ExitCode.COMPILATION_ERROR
                 }
-                script
+
+                if (!scriptFile.exists()) return invalidScript("Script file not found: $scriptFile")
+
+                if (scriptFile.isDirectory) return invalidScript("Script argument points to a directory: $scriptFile")
+
+                var script = scriptFile.toScriptSource().takeIf {
+                    scriptDefinitionProvider.isScript(it)
+                }
+                if (script == null && defaultScriptExtension != null) {
+                    script = ExplicitlyNamedFileScriptSource(
+                        scriptFile.nameWithoutExtension + defaultScriptExtension, scriptFile
+                    ).takeIf {
+                        scriptDefinitionProvider.isScript(it)
+                    }
+                }
+                script ?: return invalidScript("Unrecognized script type: ${scriptFile.name}")
             }
             else -> {
                 messageCollector.report(
@@ -157,3 +172,12 @@ fun ScriptDiagnostic.Severity.toCompilerMessageSeverity(): CompilerMessageSeveri
         ScriptDiagnostic.Severity.DEBUG -> CompilerMessageSeverity.LOGGING
     }
 
+open class ExplicitlyNamedFileScriptSource(
+    override val name: String, file: File, preloadedText: String? = null
+) : FileScriptSource(file, preloadedText), Serializable {
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}
diff --git a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/errorReporting.kt b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/errorReporting.kt
index 0e437fbb511..99f7030b227 100644
--- a/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/errorReporting.kt
+++ b/plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/impl/errorReporting.kt
@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.scripting.compiler.plugin.impl
 
 import org.jetbrains.kotlin.cli.common.arguments.Argument
 import org.jetbrains.kotlin.cli.common.arguments.K2JVMCompilerArguments
-import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.scripting.definitions.MessageReporter
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
@@ -113,6 +113,7 @@ internal fun reportArgumentsIgnoredGenerally(
         K2JVMCompilerArguments::scriptTemplates,
         K2JVMCompilerArguments::scriptResolverEnvironment,
         K2JVMCompilerArguments::disableStandardScript,
+        K2JVMCompilerArguments::defaultScriptExtension,
         K2JVMCompilerArguments::disableDefaultScriptingPlugin,
         K2JVMCompilerArguments::pluginClasspaths,
         K2JVMCompilerArguments::useJavac,
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 165a629cca6..1511162d36c 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
@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.scripting.compiler.test.linesSplitTrim
 import org.junit.Assert
 import org.junit.Test
 import java.io.File
+import java.nio.file.Files
 
 class ScriptingWithCliCompilerTest {
 
@@ -63,6 +64,85 @@ class ScriptingWithCliCompilerTest {
         )
     }
 
+    @Test
+    fun testExpressionAsMainKts() {
+        // testing that without specifying default to .main.kts, the annotation is unresolved
+        runWithK2JVMCompiler(
+            arrayOf(
+                "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
+                "-expression",
+                "\\@file:CompilerOptions(\"-Xunknown1\")"
+            ),
+            listOf(""),
+            expectedExitCode = 1,
+            expectedSomeErrPatterns = listOf(
+                "unresolved reference: CompilerOptions"
+            )
+        )
+        // it seems not possible to make a one-liner with the annotation, and
+        // annotation is the easiest available distinguishing factor for the .main.kts script
+        // so, considering "expecting an element" error as a success here
+        runWithK2JVMCompiler(
+            arrayOf(
+                "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
+                "-Xdefault-script-extension=.main.kts",
+                "-expression",
+                "\\@file:CompilerOptions(\"-Xunknown1\")"
+            ),
+            listOf(""),
+            expectedExitCode = 1,
+            expectedSomeErrPatterns = listOf(
+                "expecting an element"
+            )
+        )
+    }
+
+    @Test
+    fun testScriptAsMainKts() {
+        val scriptFile = Files.createTempFile("someScript", "").toFile()
+        scriptFile.writeText("@file:CompilerOptions(\"-abracadabra\")\n42")
+
+        // testing that without specifying default to .main.kts the script with extension .txt is not recognized
+        runWithK2JVMCompiler(
+            arrayOf(
+                "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
+                "-script",
+                scriptFile.path
+            ),
+            listOf(""),
+            expectedExitCode = 1,
+            expectedSomeErrPatterns = listOf(
+                "unrecognized script type: someScript.+"
+            )
+        )
+        runWithK2JVMCompiler(
+            arrayOf(
+                "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
+                "-Xdefault-script-extension=.main.kts",
+                "-script",
+                scriptFile.path
+            ),
+            listOf(""),
+            expectedExitCode = 1,
+            expectedSomeErrPatterns = listOf(
+                "error: invalid argument: -abracadabra"
+            )
+        )
+        runWithK2JVMCompiler(
+            arrayOf(
+                "-cp", getMainKtsClassPath().joinToString(File.pathSeparator),
+                "-Xdefault-script-extension=main.kts",
+                "-script",
+                scriptFile.path
+            ),
+            listOf(""),
+            expectedExitCode = 1,
+            expectedSomeErrPatterns = listOf(
+                "error: invalid argument: -abracadabra"
+            )
+        )
+    }
+
     @Test
     fun testExpressionWithComma() {
         runWithK2JVMCompiler(
diff --git a/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/testUtil.kt b/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/testUtil.kt
index 5e71f4fb518..6589402c4d9 100644
--- a/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/testUtil.kt
+++ b/plugins/scripting/scripting-compiler/tests/org/jetbrains/kotlin/scripting/compiler/plugin/testUtil.kt
@@ -157,8 +157,9 @@ fun runWithK2JVMCompiler(
 
 fun runWithK2JVMCompiler(
     args: Array<String>,
-    expectedOutPatterns: List<String> = emptyList(),
-    expectedExitCode: Int = 0
+    expectedAllOutPatterns: List<String> = emptyList(),
+    expectedExitCode: Int = 0,
+    expectedSomeErrPatterns: List<String>? = null
 ) {
     val (out, err, ret) = captureOutErrRet {
         CLITool.doMainNoExit(
@@ -169,15 +170,25 @@ fun runWithK2JVMCompiler(
     try {
         val outLines = out.lines()
         Assert.assertEquals(
-            "Expecting pattern:\n  ${expectedOutPatterns.joinToString("\n  ")}\nGot:\n  ${outLines.joinToString("\n  ")}",
-            expectedOutPatterns.size, outLines.size
+            "Expecting pattern:\n  ${expectedAllOutPatterns.joinToString("\n  ")}\nGot:\n  ${outLines.joinToString("\n  ")}",
+            expectedAllOutPatterns.size, outLines.size
         )
-        for ((expectedPattern, actualLine) in expectedOutPatterns.zip(outLines)) {
+        for ((expectedPattern, actualLine) in expectedAllOutPatterns.zip(outLines)) {
             Assert.assertTrue(
                 "line \"$actualLine\" do not match with expected pattern \"$expectedPattern\"",
                 Regex(expectedPattern).matches(actualLine)
             )
         }
+        if (expectedSomeErrPatterns != null) {
+            val errLines = err.lines()
+            for (expectedPattern in expectedSomeErrPatterns) {
+                val re = Regex(expectedPattern)
+                Assert.assertTrue(
+                    "Expected pattern \"$expectedPattern\" is not found in the stderr:\n${errLines.joinToString("\n")}",
+                    errLines.any { re.find(it) != null }
+                )
+            }
+        }
         Assert.assertEquals(expectedExitCode, ret.code)
     } catch (e: Throwable) {
         println("OUT:\n$out")
@@ -186,7 +197,6 @@ fun runWithK2JVMCompiler(
     }
 }
 
-
 internal fun <T> captureOutErrRet(body: () -> T): Triple<String, String, T> {
     val outStream = ByteArrayOutputStream()
     val errStream = ByteArrayOutputStream()