Rewrite handing and caching of script configurations

The refactoring could enable lazy update of the index on script loading, but seems that makeRootsChange
call doesn't work as expected, therefore initial project scan is performed in the KotlinScriptConfigurationManager
anyway

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

@@ -136,7 +136,7 @@ class KotlinCoreEnvironment private constructor(
         })
 
         KotlinScriptDefinitionProvider.getInstance(project).let { scriptDefinitionProvider ->
-            scriptDefinitionProvider.scriptDefinitions =
+            scriptDefinitionProvider.setScriptDefinitions(
                     configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS)
                             .ifEmpty {
                                 if (configuration.get(JVMConfigurationKeys.LOAD_SCRIPT_CONFIGS) ?: false)
@@ -146,14 +146,14 @@ class KotlinCoreEnvironment private constructor(
                                     }
                                 else null
                                 ?: listOf(StandardScriptDefinition)
-                            }
+                            })
 
-            configuration.addJvmClasspathRoots(
+            KotlinScriptExtraImportsProvider.getInstance(project)?.run {
-                    sourceFiles
+                configuration.addJvmClasspathRoots(
-                            .mapNotNull { src -> src.virtualFile }
+                        getCombinedClasspathFor(sourceFiles.mapNotNull { src -> src.virtualFile })
-                            .flatMap { getScriptCombinedClasspath(it, project) }
+                                .map { File(it).canonicalFile }
-                            .map { File(it).canonicalFile }
+                                .distinct())
-                            .distinct())
+            }
         }
 
         fillClasspath(configuration)
@@ -409,8 +409,9 @@ class KotlinCoreEnvironment private constructor(
         // made public for Upsource
         @JvmStatic fun registerProjectServices(projectEnvironment: JavaCoreProjectEnvironment) {
             with (projectEnvironment.project) {
-                registerService(KotlinScriptDefinitionProvider::class.java, KotlinScriptDefinitionProvider())
+                val kotlinScriptDefinitionProvider = KotlinScriptDefinitionProvider()
-                registerService(KotlinScriptExtraImportsProvider::class.java, KotlinScriptExtraImportsProvider(projectEnvironment.project))
+                registerService(KotlinScriptDefinitionProvider::class.java, kotlinScriptDefinitionProvider)
+                registerService(KotlinScriptExtraImportsProvider::class.java, KotlinScriptExtraImportsProvider(projectEnvironment.project, kotlinScriptDefinitionProvider))
                 registerService(KotlinJavaPsiFacade::class.java, KotlinJavaPsiFacade(this))
                 registerService(KtLightClassForFacade.FacadeStubCache::class.java, KtLightClassForFacade.FacadeStubCache(this))
             }

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

@@ -36,7 +36,7 @@ import org.jetbrains.kotlin.utils.PathUtil
 import java.io.File
 
 
-class KotlinConfigurableScriptDefinition(val config: KotlinScriptConfig, val environmentVars: Map<String, List<String>>?) : KotlinScriptDefinition {
+data class KotlinConfigurableScriptDefinition(val config: KotlinScriptConfig, val environmentVars: Map<String, List<String>>?) : KotlinScriptDefinition {
     override val name = config.name
     override fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter> =
         config.parameters.map { ScriptParameter(Name.identifier(it.name), getKotlinTypeByFqName(scriptDescriptor, it.type)) }

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

@@ -43,7 +43,7 @@ interface KotlinScriptDefinition {
     fun getScriptDependenciesClasspath(): List<String> = emptyList()
 }
 
-class ScriptParameter(val name: Name, val type: KotlinType)
+data class ScriptParameter(val name: Name, val type: KotlinType)
 
 object StandardScriptDefinition : KotlinScriptDefinition {
     private val ARGS_NAME = Name.identifier("args")

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

@@ -19,23 +19,34 @@ package org.jetbrains.kotlin.script
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import java.util.*
 import kotlin.concurrent.read
 import kotlin.concurrent.write
 
 class KotlinScriptDefinitionProvider {
 
     private val definitions: MutableList<KotlinScriptDefinition> = arrayListOf(StandardScriptDefinition)
-
     private val definitionsLock = java.util.concurrent.locks.ReentrantReadWriteLock()
+    private val notificationHandlers = ArrayList<() -> Unit>()
+    private val handlersLock = java.util.concurrent.locks.ReentrantReadWriteLock()
 
-    var scriptDefinitions: List<KotlinScriptDefinition>
+    fun setScriptDefinitions(newDefinitions: List<KotlinScriptDefinition>): Unit {
-        get() = definitionsLock.read { definitions } // TODO: remove as unsafe with locking, replace with particular data extractors
+        var changed = false
-        set(definitions: List<KotlinScriptDefinition>) {
+        definitionsLock.read {
-            definitionsLock.write {
+            if (newDefinitions != definitions) {
-                this.definitions.clear()
+                definitionsLock.write {
-                this.definitions.addAll(definitions)
+                    definitions.clear()
+                    definitions.addAll(newDefinitions)
+                }
+                changed = true
             }
         }
+        if (changed) {
+            handlersLock.read {
+                notificationHandlers.forEach { it() }
+            }
+        }
+    }
 
     fun findScriptDefinition(file: VirtualFile?): KotlinScriptDefinition? = definitionsLock.read {
         file?.let { file -> definitions.firstOrNull { it.isScript(file) } }
@@ -43,6 +54,10 @@ class KotlinScriptDefinitionProvider {
 
     fun isScript(file: VirtualFile?): Boolean = findScriptDefinition(file) != null
 
+    fun subscribeOnDefinitionsChanged(handler: () -> Unit): Unit {
+        handlersLock.write { notificationHandlers.add(handler) }
+    }
+
     fun addScriptDefinition(scriptDefinition: KotlinScriptDefinition) {
         definitionsLock.write {
             definitions.add(0, scriptDefinition)

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

@@ -19,34 +19,92 @@ package org.jetbrains.kotlin.script
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import java.lang.ref.WeakReference
+import java.util.*
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.read
 import kotlin.concurrent.write
 
-class KotlinScriptExtraImportsProvider(val project: Project) {
+class KotlinScriptExtraImportsProvider(val project: Project, private val scriptDefinitionProvider: KotlinScriptDefinitionProvider) {
-    private val lock = ReentrantReadWriteLock()
+    private val cacheLock = ReentrantReadWriteLock()
     private val cache = hashMapOf<VirtualFile, List<KotlinScriptExtraImport>>()
     private val envVars: Map<String, List<String>> by lazy { generateKotlinScriptClasspathEnvVars(project) }
+    private val notificationHandlers = ArrayList<(Iterable<VirtualFile>) -> Unit>()
+    private val handlersLock = java.util.concurrent.locks.ReentrantReadWriteLock()
 
-    fun isExtraImportsConfig(file: VirtualFile): Boolean = file.name.endsWith(IMPORTSFILEEXTENSION)
+    init {
-
+        val weakThis = WeakReference(this)
-    fun getExtraImports(file: VirtualFile): List<KotlinScriptExtraImport> = lock.read {
+        scriptDefinitionProvider.subscribeOnDefinitionsChanged { weakThis.get()?.invalidateAllExtraImports() }
-        if (file.isValid)
-            cache.getOrPut(file) {
-                file.parent.findFileByRelativePath(file.name + IMPORTSFILEEXTENSION)?.let {
-                    loadScriptExtraImportConfigs(it.inputStream).map { KotlinScriptExtraImportFromConfig(it, envVars) }
-                }
-                ?: emptyList()
-            }
-        else emptyList()
     }
 
-    fun invalidateExtraImports(importsFile: VirtualFile) {
+    fun isExtraImportsConfig(file: VirtualFile): Boolean = file.name.endsWith(IMPORTS_FILE_EXTENSION)
-        importsFile.parent.findFileByRelativePath(importsFile.name.removeSuffix(IMPORTSFILEEXTENSION))?.let {
+
-            lock.write {
+    fun getExtraImports(vararg files: VirtualFile): List<KotlinScriptExtraImport> = getExtraImports(files.asIterable())
-                cache.remove(it)
+
+    fun getExtraImports(files: Iterable<VirtualFile>): List<KotlinScriptExtraImport> {
+        val newCashedFiles = ArrayList<VirtualFile>()
+        val res = cacheLock.read {
+            files.flatMap { file ->
+                if (file.isValid && !file.isDirectory) {
+                    cache[file]
+                    ?: scriptDefinitionProvider.findScriptDefinition(file)?.let { def ->
+                        (listOf(KotlinScriptExtraImportFromDefinition(def)) +
+                         (file.parent.findFileByRelativePath(file.name + IMPORTS_FILE_EXTENSION)?.let {
+                             loadScriptExtraImportConfigs(it.inputStream).map { KotlinScriptExtraImportFromConfig(it, envVars) }
+                         } ?: emptyList()))
+                                .apply {
+                                    cacheLock.write { cache.put(file, this) }
+                                    newCashedFiles.add(file)
+                                }
+                    }
+                    ?: emptyList()
+                }
+                else emptyList()
             }
         }
+        notifyIfAny(newCashedFiles)
+        return res
+    }
+
+    fun invalidateExtraImportsByImportsFiles(importsFiles: Iterable<VirtualFile>) {
+        importsFiles.mapNotNull { it.parent.findFileByRelativePath(it.name.removeSuffix(IMPORTS_FILE_EXTENSION))?.let { file ->
+            cacheLock.write {
+                 cache.remove(it)?.let { file }
+            }
+        } }.let {
+            notifyIfAny(it)
+        }
+    }
+
+    fun invalidateAllExtraImports() {
+        cacheLock.write {
+            cache.keys.toList().apply {
+                cache.clear()
+            }
+        }.let {
+            notifyIfAny(it)
+        }
+    }
+
+    private fun notifyIfAny(files: Iterable<VirtualFile>) {
+        if (files.any()) {
+            handlersLock.read {
+                notificationHandlers.forEach { it(files) }
+            }
+        }
+    }
+
+    fun getKnownCombinedClasspath(): List<String> = cacheLock.read {
+        cache.values.flatMap { it.flatMap { it.classpath } }
+    }.distinct()
+
+    fun getCombinedClasspathFor(files: Iterable<VirtualFile>): List<String> =
+        getExtraImports(files)
+                .flatMap { it.classpath }
+                .distinct()
+
+    fun subscribeOnExtraImportsChanged(handler: (Iterable<VirtualFile>) -> Unit): Unit {
+        handlersLock.write { notificationHandlers.add(handler) }
     }
 
     companion object {
@@ -54,6 +112,6 @@ class KotlinScriptExtraImportsProvider(val project: Project) {
         fun getInstance(project: Project): KotlinScriptExtraImportsProvider? =
                 ServiceManager.getService(project, KotlinScriptExtraImportsProvider::class.java)
 
-        val IMPORTSFILEEXTENSION = ".ktsimports.xml"
+        val IMPORTS_FILE_EXTENSION = ".ktsimports.xml"
     }
 }

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

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

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

@@ -68,11 +68,13 @@ fun generateSampleScriptConfig(): String {
 }
 
 @Tag("scriptParam")
-class KotlinScriptParameterConfig(@Tag("name") var name: String = "",
+data class KotlinScriptParameterConfig(
-                                  @Tag("type") var type: String = "")
+        @Tag("name") var name: String = "",
+        @Tag("type") var type: String = ""
+)
 
 @Tag("script")
-class KotlinScriptConfig(
+data class KotlinScriptConfig(
         @Tag("name")
         var name: String = "KotlinScript",
 

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

@@ -61,3 +61,7 @@ class KotlinScriptExtraImportFromConfig(val config : KotlinScriptExtraImportConf
         get() = config.names
 }
 
+class KotlinScriptExtraImportFromDefinition(val scriptDefinition: KotlinScriptDefinition) : KotlinScriptExtraImport {
+    override val classpath: List<String> get() = scriptDefinition.getScriptDependenciesClasspath()
+    override val names: List<String> = emptyList()
+}

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

@@ -17,11 +17,12 @@
 package org.jetbrains.kotlin.idea.script
 
 import com.intellij.openapi.components.AbstractProjectComponent
-import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ex.ProjectRootManagerEx
 import com.intellij.openapi.util.EmptyRunnable
-import com.intellij.openapi.vfs.*
+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
 import com.intellij.psi.search.GlobalSearchScope
@@ -31,11 +32,13 @@ import org.jetbrains.kotlin.idea.caches.resolve.FileLibraryScope
 import org.jetbrains.kotlin.script.*
 import java.io.File
 import java.io.FileNotFoundException
-import java.io.InputStream
+import java.lang.ref.WeakReference
-import java.io.OutputStream
+import java.util.*
+import kotlin.concurrent.read
+import kotlin.concurrent.write
 
 @Suppress("unused") // project component
-class KotlinScriptConfigurationManager(private val project: Project,
+class KotlinScriptConfigurationManager(project: Project,
                                        private val scriptDefinitionProvider: KotlinScriptDefinitionProvider,
                                        private val scriptExtraImportsProvider: KotlinScriptExtraImportsProvider?,
                                        private val kotlinScriptDependenciesIndexableSetContributor: KotlinScriptDependenciesIndexableSetContributor?
@@ -45,47 +48,91 @@ class KotlinScriptConfigurationManager(private val project: Project,
 
     init {
         reloadScriptDefinitions()
-        val conn = myProject.messageBus.connect()
+
-        conn.subscribe(VirtualFileManager.VFS_CHANGES, object : BulkFileListener.Adapter() {
+        // TODO: get rid of this expensive call as soon as makeRootsChange call will work reliably
+        cacheAllScriptsExtraImports()
+
+        val weakThis = WeakReference(this)
+        myProject.messageBus.connect().subscribe(VirtualFileManager.VFS_CHANGES, object : BulkFileListener.Adapter() {
             override fun after(events: List<VFileEvent>) {
-                var anyScriptExtraImportsChanged = false
+                val changedExtraImportConfigs = ArrayList<VirtualFile>()
                 var anyScriptDefinitionChanged = false
                 events.filter { it is VFileEvent }.forEach {
                     it.file?.let {
-                        if (!anyScriptDefinitionChanged && isScriptDefinitionConfigFile(it)) {
+                        if (isScriptDefinitionConfigFile(it)) {
                             anyScriptDefinitionChanged = true
                         }
-                        scriptExtraImportsProvider?.run {
+                        else {
-                            if (isExtraImportsConfig(it)) {
+                            weakThis.get()?.scriptExtraImportsProvider?.run {
-                                invalidateExtraImports(it)
+                                if (isExtraImportsConfig(it)) {
-                                if (!anyScriptExtraImportsChanged) {
+                                    changedExtraImportConfigs.add(it)
-                                    anyScriptExtraImportsChanged = true
                                 }
                             }
                         }
                     }
                 }
                 if (anyScriptDefinitionChanged) {
-                    reloadScriptDefinitions()
+                    weakThis.get()?.reloadScriptDefinitions()
                 }
-                if (anyScriptExtraImportsChanged) {
+                if (changedExtraImportConfigs.isNotEmpty()) {
-                    ProjectRootManagerEx.getInstanceEx(project)?.makeRootsChange(EmptyRunnable.getInstance(), false, true)
+                    weakThis.get()?.scriptExtraImportsProvider?.invalidateExtraImportsByImportsFiles(changedExtraImportConfigs)
                 }
             }
         })
+        // omitting case then scriptExtraImportsProvider is not configured, considering it happens only in tests
+        scriptExtraImportsProvider?.subscribeOnExtraImportsChanged { files ->
+            weakThis.get()?.apply {
+                cacheLock.write {
+                    allScriptsClasspathCache = null
+                }
+                ProjectRootManagerEx.getInstanceEx(myProject)?.makeRootsChange(EmptyRunnable.getInstance(), false, true)
+            }
+        }
     }
 
-    private val scriptClasspathCache = hashMapOf<VirtualFile, List<VirtualFile>>()
     private var allScriptsClasspathCache: List<VirtualFile>? = null
+    private val cacheLock = java.util.concurrent.locks.ReentrantReadWriteLock()
 
     fun getScriptClasspath(file: VirtualFile): List<VirtualFile> =
-            scriptClasspathCache.getOrPut(file, {
+            scriptExtraImportsProvider
-                getScriptClasspathRaw(file)
+                    ?.getExtraImports(file)
-                        .mapNotNull { StandardFileSystems.local().findFileByPath(it) }
+                    ?.flatMap { it.classpath }
-                        .distinct()
+                    ?.map { it.classpathEntryToVfs() }
-            })
+                    ?: emptyList()
 
-    fun getAllScriptsClasspath(): List<VirtualFile> {
+    fun getAllScriptsClasspath(): List<VirtualFile> = cacheLock.read {
+        if (allScriptsClasspathCache == null) {
+            allScriptsClasspathCache =
+                    (scriptExtraImportsProvider?.getKnownCombinedClasspath() ?: emptyList())
+                    .distinct()
+                    .mapNotNull { it.classpathEntryToVfs() }
+        }
+        return allScriptsClasspathCache!!
+    }
+
+    private fun String.classpathEntryToVfs(): VirtualFile =
+            if (File(this).isDirectory)
+                StandardFileSystems.local()?.findFileByPath(this) ?: throw FileNotFoundException("Classpath entry points to a non-existent location: ${this}")
+            else
+                StandardFileSystems.jar()?.findFileByPath(this + URLUtil.JAR_SEPARATOR) ?: throw FileNotFoundException("Classpath entry points to a file that is not a JAR archive: ${this}")
+
+    fun getAllScriptsClasspathScope(): GlobalSearchScope? {
+        return getAllScriptsClasspath().let { cp ->
+            if (cp.isEmpty()) null
+            else GlobalSearchScope.union(cp.map { FileLibraryScope(myProject, it) }.toTypedArray())
+        }
+    }
+
+    private fun reloadScriptDefinitions() {
+        loadScriptConfigsFromProjectRoot(File(myProject.basePath ?: ".")).let {
+            if (it.isNotEmpty()) {
+                scriptDefinitionProvider.setScriptDefinitions(
+                        it.map { KotlinConfigurableScriptDefinition(it, kotlinEnvVars) } + StandardScriptDefinition)
+            }
+        }
+    }
+
+    private fun cacheAllScriptsExtraImports() {
         fun<R> VirtualFile.vfsWalkFiles(onFile: (VirtualFile) -> List<R>?): List<R> {
             assert(isDirectory)
             return children.flatMap { when {
@@ -93,38 +140,8 @@ class KotlinScriptConfigurationManager(private val project: Project,
                 else -> onFile(it) ?: emptyList()
             } }
         }
-        if (allScriptsClasspathCache == null) {
+        myProject.baseDir.vfsWalkFiles {
-            allScriptsClasspathCache = project.baseDir.vfsWalkFiles { getScriptClasspathRaw(it) }
+            scriptExtraImportsProvider?.getExtraImports(it)
-                    .distinct()
-                    .mapNotNull {
-                        if (File(it).isDirectory)
-                            StandardFileSystems.local()?.findFileByPath(it) ?: throw FileNotFoundException("Classpath entry points to a non-existent location: $it")
-                        else
-                            StandardFileSystems.jar()?.findFileByPath(it + URLUtil.JAR_SEPARATOR) ?: throw FileNotFoundException("Classpath entry points to a file that is not a JAR archive: $it")
-
-                    }
-        }
-        return allScriptsClasspathCache!!
-    }
-
-    fun getAllScriptsClasspathScope(): GlobalSearchScope? {
-        return getAllScriptsClasspath().let { cp ->
-            if (cp.isEmpty()) null
-            else GlobalSearchScope.union(cp.map { FileLibraryScope(project, it) }.toTypedArray())
-        }
-    }
-
-    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()) {
-                scriptDefinitionProvider.scriptDefinitions =
-                        it.map { KotlinConfigurableScriptDefinition(it, kotlinEnvVars) } + StandardScriptDefinition
-            }
         }
     }