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

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) {