Implement class finder for customized scripts resolving - not yet completely functional

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

@@ -226,7 +226,8 @@ object KotlinToJVMBytecodeCompiler {
         try {
             try {
                 tryConstructClass(scriptClass.kotlin, scriptArgs)
-                    ?: throw RuntimeException("unable to find appropriate constructor for class ${scriptClass.name} accepting arguments $scriptArgs")
+                    ?: throw RuntimeException("unable to find appropriate constructor for class ${scriptClass.name} accepting arguments $scriptArgs\n" +
+                                "\tconstructors: \n\t\t${scriptClass.kotlin.constructors.joinToString("\n\t\t", "(") { it.parameters.joinToString { it.type.toString() } }}")
             }
             finally {
                 // NB: these lines are required (see KT-9546) but aren't covered by tests
@@ -291,16 +292,19 @@ object KotlinToJVMBytecodeCompiler {
                     else -> null
                 }
 
-        fun foldingFunc(state: Pair<List<Any>, List<String>>, par: KParameter): Pair<List<Any>, List<String>> {
+        fun foldingFunc(state: Pair<List<Any>, List<String>?>, par: KParameter): Pair<List<Any>, List<String>?> {
-            if (state.second.isNotEmpty()) {
+            state.second?.let { scriptArgsLeft ->
                 try {
-                    val primArgCandidate = convertPrimitive(par.type, state.second.first())
+                    if (scriptArgsLeft.isNotEmpty()) {
-                    if (primArgCandidate != null)
+                        val primArgCandidate = convertPrimitive(par.type, scriptArgsLeft.first())
-                        return Pair(state.first + primArgCandidate, state.second.drop(1))
+                        if (primArgCandidate != null)
+                            return@foldingFunc Pair(state.first + primArgCandidate, scriptArgsLeft.drop(1))
+                    }
 
-                    val arrayArgCandidate = convertArray((par.type.javaType as? Class<*>)?.componentType?.kotlin?.defaultType, state.second)
+                    val arrCompType = (par.type.javaType as? Class<*>)?.componentType?.kotlin?.defaultType
+                    val arrayArgCandidate = convertArray(arrCompType, scriptArgsLeft)
                     if (arrayArgCandidate != null)
-                        return Pair(state.first + arrayArgCandidate, emptyList<String>())
+                        return@foldingFunc Pair(state.first + arrayArgCandidate, null)
                 }
                 catch (e: NumberFormatException) {
                 } // just skips to return below
@@ -310,7 +314,7 @@ object KotlinToJVMBytecodeCompiler {
 
         for (ctor in scriptClass.constructors) {
             val (ctorArgs, scriptArgsLeft) = ctor.parameters.fold(Pair(emptyList<Any>(), scriptArgs), ::foldingFunc)
-            if (ctorArgs.size == ctor.parameters.size && scriptArgsLeft.isEmpty())
+            if (ctorArgs.size == ctor.parameters.size && (scriptArgsLeft == null || scriptArgsLeft.isEmpty()))
                 return ctor.call(*ctorArgs.toTypedArray())
         }
         return null

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

@@ -50,20 +50,16 @@ object StandardScriptDefinition : KotlinScriptDefinition {
 
     override val name = "Kotlin Script"
 
-    override fun getScriptName(script: KtScript): Name {
+    override fun getScriptName(script: KtScript): Name =
-        return ScriptNameUtil.fileNameWithExtensionStripped(script, KotlinParserDefinition.STD_SCRIPT_EXT)
+            ScriptNameUtil.fileNameWithExtensionStripped(script, KotlinParserDefinition.STD_SCRIPT_EXT)
-    }
 
-    override fun isScript(file: VirtualFile): Boolean {
+    override fun isScript(file: VirtualFile): Boolean =
-        return PathUtil.getFileExtension(file.name) == KotlinParserDefinition.STD_SCRIPT_SUFFIX
+        PathUtil.getFileExtension(file.name) == 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
-    override fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter> {
+    override fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter> =
-        val kc: KClass<StandardScriptDefinition> = StandardScriptDefinition::class
+            makeStringListScriptParameters(scriptDescriptor, ARGS_NAME)
-        return makeStringListScriptParameters(scriptDescriptor, ARGS_NAME)
-    }
 }
 
 fun makeStringListScriptParameters(scriptDescriptor: ScriptDescriptor, propertyName: Name): List<ScriptParameter> {

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

@@ -33,3 +33,11 @@ fun getScriptExtraImports(psiFile: PsiFile): List<KotlinScriptExtraImport>  =
         psiFile.virtualFile?.let { file ->
             KotlinScriptExtraImportsProvider.getInstance(psiFile.project)?.getExtraImports(file)
         } ?: emptyList()
+
+fun getScriptCombinedClasspath(file: VirtualFile, project: Project): List<String> =
+        (getScriptDefinition(file, project)?.getScriptDependenciesClasspath() ?: emptyList()) +
+        getScriptExtraImports(file, project).flatMap { it.classpath }
+
+fun getScriptCombinedClasspath(psiFile: PsiFile): List<String> =
+        (getScriptDefinition(psiFile) ?.getScriptDependenciesClasspath() ?: emptyList()) +
+        getScriptExtraImports(psiFile).flatMap { it.classpath }

compiler/testData/script/fib_std.kts

@@ -1,4 +1,4 @@
-// Expecting two string parameters
+// Expecting two string parameters or nothing
 
 fun fib(n: Int): Int {
     val v = if(n < 2) 1 else fib(n-1) + fib(n-2)
@@ -6,5 +6,8 @@ fun fib(n: Int): Int {
     return v
 }
 
-System.out.println("num: ${args[0]} (${args[1]})")
+val num = if (args.size > 0) java.lang.Integer.parseInt(args[0]) else 4
-val result = fib(java.lang.Integer.parseInt(args[0]))
+val comment = if (args.size > 1) args[1] else "none"
+
+System.out.println("num: $num ($comment)")
+val result = fib(num)

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

@@ -67,6 +67,13 @@ class ReflectedSuperclassWithParamsTestScriptDefinition(extension: String,
             superclassParameters.map { it.name }
 }
 
+class StandardWithClasspathScriptDefinition(extension: String, classpath: List<String>? = null)
+    : BaseScriptDefinition(extension, classpath)
+{
+    override fun getScriptParameters(scriptDescriptor: ScriptDescriptor) =
+            StandardScriptDefinition.getScriptParameters(scriptDescriptor)
+}
+
 fun classpathFromProperty(): List<String> =
     System.getProperty("java.class.path")?.let {
         it.split(String.format("\\%s", File.pathSeparatorChar).toRegex()).dropLastWhile { it.isEmpty() }.toTypedArray()

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

@@ -39,11 +39,10 @@ import org.jetbrains.kotlin.utils.PathUtil
 import org.junit.Assert
 import org.junit.Test
 import java.io.File
-import kotlin.reflect.defaultType
+import java.lang.reflect.InvocationTargetException
 
 class ScriptTest {
     @Test
-    @Throws(Exception::class)
     fun testScriptWithParam() {
         val aClass = compileScript("fib.kts", SimpleParamsTestScriptDefinition(".kts", numIntParam()))
         Assert.assertNotNull(aClass)
@@ -51,8 +50,7 @@ class ScriptTest {
     }
 
     @Test
-    @Throws(Exception::class)
+    fun testStandardScriptWithParams() {
-    fun testStandardScript() {
         val aClass = compileScript("fib_std.kts", StandardScriptDefinition)
         Assert.assertNotNull(aClass)
         val anObj = KotlinToJVMBytecodeCompiler.tryConstructClassPub(aClass!!.kotlin, listOf("4", "comment"))
@@ -60,7 +58,14 @@ class ScriptTest {
     }
 
     @Test
-    @Throws(Exception::class)
+    fun testStandardScriptWithoutParams() {
+        val aClass = compileScript("fib_std.kts", StandardScriptDefinition)
+        Assert.assertNotNull(aClass)
+        val anObj = KotlinToJVMBytecodeCompiler.tryConstructClassPub(aClass!!.kotlin, emptyList())
+        Assert.assertNotNull(anObj)
+    }
+
+    @Test
     fun testScriptWithParamConversion() {
         val aClass = compileScript("fib.kts", SimpleParamsTestScriptDefinition(".kts", numIntParam()))
         Assert.assertNotNull(aClass)
@@ -69,7 +74,6 @@ class ScriptTest {
     }
 
     @Test
-    @Throws(Exception::class)
     fun testScriptWithPackage() {
         val aClass = compileScript("fib.pkg.kts", SimpleParamsTestScriptDefinition(".kts", numIntParam()))
         Assert.assertNotNull(aClass)
@@ -77,7 +81,6 @@ class ScriptTest {
     }
 
     @Test
-    @Throws(Exception::class)
     fun testScriptWithScriptDefinition() {
         val aClass = compileScript("fib.fib.kt", SimpleParamsTestScriptDefinition(".fib.kt", numIntParam()))
         Assert.assertNotNull(aClass)
@@ -124,11 +127,45 @@ class ScriptTest {
         Assert.assertNotNull(aClass2)
     }
 
+    @Test
+    fun testSmokeScriptException() {
+        val aClass = compileSmokeTestScript(
+                "scriptException/script.kts",
+                StandardWithClasspathScriptDefinition(
+                        ".kts",
+                        listOf("dependencies/bootstrap-compiler/Kotlin/kotlinc/lib/kotlin-runtime.jar",
+                               "dependencies/bootstrap-compiler/Kotlin/kotlinc/lib/kotlin-reflect.jar")))
+        Assert.assertNotNull(aClass)
+        var exceptionThrown = false
+        try {
+            KotlinToJVMBytecodeCompiler.tryConstructClassPub(aClass!!.kotlin, emptyList())
+        }
+        catch (e: InvocationTargetException) {
+            Assert.assertTrue(e.cause is IllegalStateException)
+            exceptionThrown = true
+        }
+        Assert.assertTrue(exceptionThrown)
+    }
+
     private fun compileScript(
             scriptPath: String,
             scriptDefinition: KotlinScriptDefinition,
             runIsolated: Boolean = true,
-            suppressOutput: Boolean = false): Class<*>?
+            suppressOutput: Boolean = false): Class<*>? =
+    compileScriptImpl("compiler/testData/script/" + scriptPath, scriptDefinition, runIsolated, suppressOutput)
+
+    private fun compileSmokeTestScript(
+            scriptPath: String,
+            scriptDefinition: KotlinScriptDefinition,
+            runIsolated: Boolean = true,
+            suppressOutput: Boolean = false): Class<*>? =
+            compileScriptImpl("compiler/testData/integration/smoke/" + scriptPath, scriptDefinition, runIsolated, suppressOutput)
+
+    private fun compileScriptImpl(
+            scriptPath: String,
+            scriptDefinition: KotlinScriptDefinition,
+            runIsolated: Boolean,
+            suppressOutput: Boolean): Class<*>?
     {
         val paths = PathUtil.getKotlinPathsForDistDirectory()
         val messageCollector =
@@ -139,7 +176,7 @@ class ScriptTest {
         try {
             val configuration = KotlinTestUtils.newConfiguration(ConfigurationKind.JDK_ONLY, TestJdkKind.FULL_JDK)
             configuration.put(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector)
-            configuration.addKotlinSourceRoot("compiler/testData/script/" + scriptPath)
+            configuration.addKotlinSourceRoot(scriptPath)
             configuration.add(JVMConfigurationKeys.SCRIPT_DEFINITIONS, scriptDefinition)
             scriptDefinition.getScriptDependenciesClasspath().forEach { configuration.addJvmClasspathRoot(File(it)) }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/IdeaModuleInfos.kt

@@ -28,6 +28,7 @@ import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
+import com.intellij.psi.search.DelegatingGlobalSearchScope
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
@@ -289,6 +290,12 @@ internal object NotUnderContentRootModuleInfo : IdeaModuleInfo {
     override fun dependencies(): List<IdeaModuleInfo> = listOf(this)
 }
 
+class CustomizedScriptModuleSearchScope(val scriptFile: VirtualFile, baseScope: GlobalSearchScope) : DelegatingGlobalSearchScope(baseScope) {
+    override fun equals(other: Any?) = other is CustomizedScriptModuleSearchScope && scriptFile == other.scriptFile && super.equals(other)
+
+    override fun hashCode() = scriptFile.hashCode() * 73 * super.hashCode()
+}
+
 internal data class CustomizedScriptModuleInfo(val project: Project, val module: Module?, val virtualFile: VirtualFile,
                                                val scriptDefinition: KotlinScriptDefinition,
                                                val scriptExtraImports: List<KotlinScriptExtraImport>) : IdeaModuleInfo {
@@ -297,7 +304,7 @@ internal data class CustomizedScriptModuleInfo(val project: Project, val module:
 
     override val name: Name = Name.special("<$SCRIPT_NAME_PREFIX${scriptDefinition.name}>")
 
-    override fun contentScope() = GlobalSearchScope.union(dependenciesRoots().map { FileLibraryScope(project, it) }.toTypedArray())
+    override fun contentScope() = CustomizedScriptModuleSearchScope(virtualFile, GlobalSearchScope.union(dependenciesRoots().map { FileLibraryScope(project, it) }.toTypedArray()))
 
     private fun dependenciesRoots(): List<VirtualFile> {
         // TODO: find out whether it should be cashed (some changes listener should be implemented for the cached roots)

idea/src/META-INF/plugin.xml

@@ -571,6 +571,7 @@
         implementationClass="org.jetbrains.kotlin.idea.hierarchy.overrides.KotlinOverrideHierarchyProvider" />
 
     <java.elementFinder implementation="org.jetbrains.kotlin.asJava.JavaElementFinder"/>
+    <java.elementFinder implementation="org.jetbrains.kotlin.idea.script.KotlinScriptDependenciesClassFinder"/>
     <java.shortNamesCache implementation="org.jetbrains.kotlin.idea.caches.KotlinShortNamesCache"/>
 
     <stubElementTypeHolder class="org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes"/>

idea/src/org/jetbrains/kotlin/idea/script/KotlinScriptConfigurationManager.kt

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.idea.script
 
 import com.intellij.openapi.components.AbstractProjectComponent
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.StandardFileSystems
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.openapi.vfs.newvfs.BulkFileListener
 import com.intellij.openapi.vfs.newvfs.events.VFileEvent
@@ -25,9 +27,9 @@ import org.jetbrains.kotlin.script.*
 import java.io.File
 
 @Suppress("unused") // project component
-class KotlinScriptConfigurationManager(project: Project,
+class KotlinScriptConfigurationManager(private val project: Project,
                                        private val scriptDefinitionProvider: KotlinScriptDefinitionProvider,
-                                       scriptExtraImportsProvider: KotlinScriptExtraImportsProvider?
+                                       private val scriptExtraImportsProvider: KotlinScriptExtraImportsProvider?
 ) : AbstractProjectComponent(project) {
 
     private val kotlinEnvVars: Map<String, List<String>> by lazy { generateKotlinScriptClasspathEnvVarsForIdea(myProject) }
@@ -57,6 +59,31 @@ class KotlinScriptConfigurationManager(project: Project,
         })
     }
 
+    // TODO: cache
+    fun getScriptClasspath(file: VirtualFile): List<VirtualFile> =
+        getScriptClasspathRaw(file)
+                .mapNotNull { StandardFileSystems.local().findFileByPath(it) }
+                .distinct()
+
+    // TODO: cache
+    fun getAllScriptsClasspath(): List<VirtualFile> {
+        fun<R> VirtualFile.vfsWalkFiles(onFile: (VirtualFile) -> List<R>?): List<R> {
+            assert(isDirectory)
+            return children.flatMap { when {
+                it.isDirectory -> it.vfsWalkFiles(onFile)
+                else -> onFile(it) ?: emptyList()
+            } }
+        }
+        return project.baseDir.vfsWalkFiles { getScriptClasspathRaw(it) }
+                .mapNotNull { StandardFileSystems.local()?.findFileByPath(it) }
+                .distinct()
+    }
+
+    private fun getScriptClasspathRaw(file: VirtualFile): List<String> =
+            scriptDefinitionProvider.findScriptDefinition(file)?.getScriptDependenciesClasspath()?.let {
+                it + (scriptExtraImportsProvider?.getExtraImports(file)?.flatMap { it.classpath } ?: emptyList())
+            } ?: emptyList()
+
     private fun reloadScriptDefinitions() {
         loadScriptConfigsFromProjectRoot(File(myProject.basePath ?: ".")).let {
             if (it.isNotEmpty()) {

idea/src/org/jetbrains/kotlin/idea/script/KotlinScriptDependenciesClassFinder.kt

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.idea.script
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.roots.ProjectFileIndex
+import com.intellij.openapi.roots.impl.PackageDirectoryCache
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.NonClasspathClassFinder
+import com.intellij.psi.PsiClass
+import com.intellij.psi.search.EverythingGlobalScope
+import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.util.containers.ConcurrentFactoryMap
+import org.jetbrains.kotlin.idea.caches.resolve.CustomizedScriptModuleSearchScope
+
+@Suppress("unused") // project extension
+class KotlinScriptDependenciesClassFinder(project: Project,
+                                         private val kotlinScriptConfigurationManager: KotlinScriptConfigurationManager
+) : NonClasspathClassFinder(project) {
+
+    private val myCaches = object : ConcurrentFactoryMap<VirtualFile, PackageDirectoryCache>() {
+        override fun create(file: VirtualFile): PackageDirectoryCache? = NonClasspathClassFinder.createCache( kotlinScriptConfigurationManager.getScriptClasspath(file))
+    }
+
+    override fun calcClassRoots(): List<VirtualFile> = kotlinScriptConfigurationManager.getAllScriptsClasspath()
+
+    override fun getCache(scope: GlobalSearchScope?): PackageDirectoryCache =
+            (scope as? CustomizedScriptModuleSearchScope)?.let { myCaches.get(it.scriptFile) } ?: super.getCache(scope)
+
+    override fun clearCache() {
+        super.clearCache()
+        myCaches.clear()
+    }
+
+    override fun findClass(qualifiedName: String, scope: GlobalSearchScope): PsiClass? =
+        super.findClass(qualifiedName, scope)?.let { aClass ->
+            when {
+                scope is CustomizedScriptModuleSearchScope ||
+                scope is EverythingGlobalScope ||
+                aClass.containingFile?.virtualFile.let { file ->
+                    file != null &&
+                    with (ProjectFileIndex.SERVICE.getInstance(myProject)) {
+                        !isInContent(file) &&
+                        !isInLibraryClasses(file) &&
+                        !isInLibrarySource(file)
+                    }
+                } -> aClass
+                else -> null
+            }
+        }
+}

idea/tests/org/jetbrains/kotlin/idea/RegisteredFindersTest.kt

@@ -27,7 +27,7 @@ class RegisteredFindersTest : KotlinLightCodeInsightFixtureTestCase() {
     override fun getProjectDescriptor() = LightCodeInsightFixtureTestCase.JAVA_LATEST
 
     fun testKnownNonClasspathFinder() {
-        val expectedFindersNames = setOf("GantClassFinder", "GradleClassFinder").toMutableSet()
+        val expectedFindersNames = setOf("GantClassFinder", "GradleClassFinder", "KotlinScriptDependenciesClassFinder").toMutableSet()
 
         project.getExtensions<PsiElementFinder>(PsiElementFinder.EP_NAME).forEach { finder ->
             if (finder is NonClasspathClassFinder) {