Make script detection API generic, fix tests

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/ReplInterpreter.kt

@@ -239,7 +239,7 @@ class ReplInterpreter(
 
             override fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter> = emptyList()
 
-            override fun isScript(file: VirtualFile): Boolean = StandardScriptDefinition.isScript(file)
+            override fun <TF> isScript(file: TF): Boolean = StandardScriptDefinition.isScript(file)
 
             override fun getScriptName(script: KtScript): Name = StandardScriptDefinition.getScriptName(script)
         }

compiler/frontend/src/org/jetbrains/kotlin/parsing/KotlinParser.java

@@ -43,7 +43,7 @@ public class KotlinParser implements PsiParser {
     @NotNull
     public ASTNode parse(IElementType iElementType, PsiBuilder psiBuilder, PsiFile psiFile) {
         KotlinParsing jetParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(psiBuilder));
-        if (scriptDefinitionProvider != null && scriptDefinitionProvider.isScript(psiFile.getOriginalFile().getVirtualFile())
+        if (scriptDefinitionProvider != null && scriptDefinitionProvider.isScript(psiFile)
             || psiFile.getName().endsWith(KotlinParserDefinition.STD_SCRIPT_EXT)) {
             jetParsing.parseScript();
         }

compiler/frontend/src/org/jetbrains/kotlin/script/KotlinConfigurableScriptDefinition.kt

@@ -47,8 +47,8 @@ data class KotlinConfigurableScriptDefinition(val config: KotlinScriptConfig, va
     override fun getScriptParametersToPassToSuperclass(scriptDescriptor: ScriptDescriptor): List<Name> =
         config.superclassParamsMapping.map { Name.identifier(it) }
 
-    override fun isScript(file: VirtualFile): Boolean =
+    override fun <TF> isScript(file: TF): Boolean =
-            Regex(config.fileNameMatch).matches(file.name)
+            Regex(config.fileNameMatch).matches(getFileName(file))
 
     override fun getScriptName(script: KtScript): Name = ScriptNameUtil.fileNameWithExtensionStripped(script, KotlinParserDefinition.STD_SCRIPT_EXT)
 

compiler/frontend/src/org/jetbrains/kotlin/script/KotlinScriptDefinition.kt

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.script
 
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiFile
 import com.intellij.util.PathUtil
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.name.ClassId
@@ -31,6 +32,7 @@ import org.jetbrains.kotlin.serialization.deserialization.findNonGenericClassAcr
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.Variance
+import java.io.File
 import kotlin.reflect.KClass
 
 interface KotlinScriptDefinition {
@@ -38,13 +40,22 @@ interface KotlinScriptDefinition {
     fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter>
     fun getScriptSupertypes(scriptDescriptor: ScriptDescriptor): List<KotlinType> = emptyList()
     fun getScriptParametersToPassToSuperclass(scriptDescriptor: ScriptDescriptor): List<Name> = emptyList()
-    fun isScript(file: VirtualFile): Boolean
+    fun <TF> isScript(file: TF): Boolean
     fun getScriptName(script: KtScript): Name
     fun getScriptDependenciesClasspath(): List<String> = emptyList()
 }
 
 data class ScriptParameter(val name: Name, val type: KotlinType)
 
+fun <TF> getFileExtension(file: TF) = PathUtil.getFileExtension(getFileName(file))
+
+fun <TF> getFileName(file: TF): String = when (file) {
+    is PsiFile -> file.originalFile.name
+    is VirtualFile -> file.name
+    is File -> file.name
+    else -> throw IllegalArgumentException("Unsupported file type $file")
+}
+
 object StandardScriptDefinition : KotlinScriptDefinition {
     private val ARGS_NAME = Name.identifier("args")
 
@@ -53,8 +64,8 @@ object StandardScriptDefinition : KotlinScriptDefinition {
     override fun getScriptName(script: KtScript): Name =
             ScriptNameUtil.fileNameWithExtensionStripped(script, KotlinParserDefinition.STD_SCRIPT_EXT)
 
-    override fun isScript(file: VirtualFile): Boolean =
+    override fun <TF> isScript(file: TF): Boolean =
-        PathUtil.getFileExtension(file.name) == KotlinParserDefinition.STD_SCRIPT_SUFFIX
+        getFileExtension(file) == KotlinParserDefinition.STD_SCRIPT_SUFFIX
 
     // NOTE: for now we treat .kts files as if they have 'args: Array<String>' parameter
     // this is not supposed to be final design

compiler/frontend/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionProvider.kt

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.script
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiFile
 import java.util.*
 import kotlin.concurrent.read
 import kotlin.concurrent.write
@@ -48,11 +49,11 @@ class KotlinScriptDefinitionProvider {
         }
     }
 
-    fun findScriptDefinition(file: VirtualFile?): KotlinScriptDefinition? = definitionsLock.read {
+    fun<TF> findScriptDefinition(file: TF): KotlinScriptDefinition? = definitionsLock.read {
-        file?.let { file -> definitions.firstOrNull { it.isScript(file) } }
+        definitions.firstOrNull { it.isScript(file) }
     }
 
-    fun isScript(file: VirtualFile?): Boolean = findScriptDefinition(file) != null
+    fun<TF> isScript(file: TF): Boolean = findScriptDefinition(file) != null
 
     fun subscribeOnDefinitionsChanged(handler: () -> Unit): Unit {
         handlersLock.write { notificationHandlers.add(handler) }

compiler/frontend/src/org/jetbrains/kotlin/script/getScriptDefinition.kt

@@ -24,7 +24,7 @@ fun getScriptDefinition(file: VirtualFile, project: Project): KotlinScriptDefini
     KotlinScriptDefinitionProvider.getInstance(project).findScriptDefinition(file)
 
 fun getScriptDefinition(psiFile: PsiFile): KotlinScriptDefinition? =
-    KotlinScriptDefinitionProvider.getInstance(psiFile.project).findScriptDefinition(psiFile.originalFile.virtualFile)
+    KotlinScriptDefinitionProvider.getInstance(psiFile.project).findScriptDefinition(psiFile)
 
 fun getScriptExtraImports(file: VirtualFile, project: Project): List<KotlinScriptExtraImport>  =
         KotlinScriptExtraImportsProvider.getInstance(project)?.getExtraImports(file) ?: emptyList()

compiler/testData/script/fib.fib.kt

@@ -1,10 +0,0 @@
-// this script expected parameter num : Int
-
-fun fib(n: Int): Int {
-    val v = if(n < 2) 1 else fib(n-1) + fib(n-2)
-    System.out.println("fib($n)=$v")
-    return v
-}
-
-System.out.println("num: $num")
-val result = fib(num)

compiler/tests-common/org/jetbrains/kotlin/scripts/TestScriptDefinitions.kt

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.scripts
 
-import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtScript
@@ -29,7 +28,7 @@ import kotlin.reflect.KClass
 
 abstract class BaseScriptDefinition (val extension: String, val classpath: List<String>? = null) : KotlinScriptDefinition {
     override val name = "Test Kotlin Script"
-    override fun isScript(file: VirtualFile): Boolean = file.name.endsWith(extension)
+    override fun <TF> isScript(file: TF): Boolean = org.jetbrains.kotlin.script.getFileName(file).endsWith(extension)
     override fun getScriptName(script: KtScript): Name = ScriptNameUtil.fileNameWithExtensionStripped(script, extension)
     override fun getScriptDependenciesClasspath(): List<String> =
             classpath ?: (classpathFromProperty() + classpathFromClassloader(BaseScriptDefinition::class.java.classLoader)).distinct()

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

@@ -83,7 +83,7 @@ class ScriptTest {
 
     @Test
     fun testScriptWithScriptDefinition() {
-        val aClass = compileScript("fib.fib.kt", SimpleParamsTestScriptDefinition(".fib.kt", numIntParam()), null)
+        val aClass = compileScript("fib.kts", SimpleParamsTestScriptDefinition(".kts", numIntParam()), null)
         Assert.assertNotNull(aClass)
         aClass!!.getConstructor(Integer.TYPE).newInstance(4)
     }
@@ -213,7 +213,6 @@ class ScriptTest {
                 MessageCollectorUtil.reportException(messageCollector, t)
                 throw t
             }
-
         }
         finally {
             Disposer.dispose(rootDisposable)

idea/idea-repl/src/org/jetbrains/kotlin/console/KotlinConsoleRunner.kt

@@ -44,6 +44,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.vfs.CharsetToolkit
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.PsiManager
 import com.intellij.psi.impl.PsiFileFactoryImpl
@@ -123,7 +124,14 @@ class KotlinConsoleRunner(
 
     private val consoleScriptDefinition = object : KotlinScriptDefinition {
         override val name = "Kotlin REPL"
-        override fun isScript(file: VirtualFile) = file == consoleView.virtualFile
+        override fun <TF> isScript(file: TF): Boolean {
+            val vf = when (file) {
+                is PsiFile -> file.originalFile.virtualFile
+                is VirtualFile -> file
+                else -> return false
+            }
+            return vf == consoleView.virtualFile
+        }
         override fun getScriptParameters(scriptDescriptor: ScriptDescriptor) = emptyList<ScriptParameter>()
         override fun getScriptName(script: KtScript) = Name.identifier("REPL")
     }