Optimize and simplify script deps cache and appropriate index rebuilding

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

@@ -28,8 +28,7 @@ import kotlin.concurrent.write
 class KotlinScriptExternalImportsProvider(val project: Project, private val scriptDefinitionProvider: KotlinScriptDefinitionProvider) {
 
     private val cacheLock = ReentrantReadWriteLock()
-    private val cache = hashMapOf<String, KotlinScriptExternalDependencies>()
+    private val cache = hashMapOf<String, KotlinScriptExternalDependencies?>()
-    private val cacheOfNulls = hashSetOf<String>()
 
     fun <TF: Any> getExternalImports(file: TF): KotlinScriptExternalDependencies? = cacheLock.read { calculateExternalDependencies(file) }
 
@@ -39,58 +38,62 @@ class KotlinScriptExternalImportsProvider(val project: Project, private val scri
 
     private fun <TF: Any> calculateExternalDependencies(file: TF): KotlinScriptExternalDependencies? {
         val path = getFilePath(file)
-        return cache[path]
+        val cached = cache[path]
-               ?: if (cacheOfNulls.contains(path)) null
+        return if (cached != null) cached else {
-               else scriptDefinitionProvider.findScriptDefinition(file)?.getDependenciesFor(file, project, null)
+            val scriptDef = scriptDefinitionProvider.findScriptDefinition(file)
-                    .apply {
+            if (scriptDef != null) {
-                        if (this != null) {
+                val deps = scriptDef.getDependenciesFor(file, project, null)
-                            log.info("[kts] new cached deps for $path: ${this.classpath.joinToString(File.pathSeparator)}")
+                if (deps != null) {
-                        }
+                    log.info("[kts] new cached deps for $path: ${deps.classpath.joinToString(File.pathSeparator)}")
-                        cacheLock.write {
+                }
-                            if (this == null) {
+                cacheLock.write {
-                                cacheOfNulls.add(path)
+                    cache.put(path, deps)
-                            }
+                }
-                            else {
+                deps
-                                cache.put(path, this)
+            }
-                            }
+            else null
-                        }
+        }
-                    }
     }
 
     // optimized for initial caching, additional handling of possible duplicates to save a call to distinct
-    fun <TF: Any> cacheExternalImports(files: Iterable<TF>): Unit = cacheLock.write {
+    // returns list of cached files
+    fun <TF: Any> cacheExternalImports(files: Iterable<TF>): Iterable<TF> = cacheLock.write {
         val uncached = hashSetOf<String>()
-        files.forEach { file ->
+        files.mapNotNull { file ->
             val path = getFilePath(file)
-            if (isValidFile(file) && !cache.containsKey(path) && !cacheOfNulls.contains(path) && !uncached.contains(path)) {
+            if (isValidFile(file) && !cache.containsKey(path) && !uncached.contains(path)) {
                 val scriptDef = scriptDefinitionProvider.findScriptDefinition(file)
                 if (scriptDef != null) {
                     val deps = scriptDef.getDependenciesFor(file, project, null)
-                    log.info("[kts] cached deps for $path: ${deps?.classpath?.joinToString(File.pathSeparator)}")
                     if (deps != null) {
-                        cache.put(path, deps)
+                        log.info("[kts] cached deps for $path: ${deps.classpath.joinToString(File.pathSeparator)}")
-                    }
-                    else {
-                        cacheOfNulls.add(path)
                     }
+                    cache.put(path, deps)
+                    file
                 }
                 else {
                     uncached.add(path)
+                    null
                 }
             }
+            else null
         }
     }
 
     // optimized for update, no special duplicates handling
+    // returns files with valid script definition (or deleted from cache - which in fact should have script def too)
+    // TODO: this is the badly designed contract, since it mixes the entities, but these files are needed on the calling site now. Find out other solution
     fun <TF: Any> updateExternalImportsCache(files: Iterable<TF>): Iterable<TF> = cacheLock.write {
         files.mapNotNull { file ->
             val path = getFilePath(file)
             if (!isValidFile(file)) {
-                if (cache.remove(path) != null || cacheOfNulls.remove(path)) {
+                if (cache.remove(path) != null) {
-                    log.info("[kts] removed deps for invalid file $path")
+                    log.debug("[kts] removed deps for file $path")
                     file
                 } // cleared
-                else null // unknown
+                else {
+                    null // unknown
+                }
             }
             else {
                 val scriptDef = scriptDefinitionProvider.findScriptDefinition(file)
@@ -103,22 +106,17 @@ class KotlinScriptExternalImportsProvider(val project: Project, private val scri
                             // changed or new
                             log.info("[kts] updated/new cached deps for $path: ${deps.classpath.joinToString(File.pathSeparator)}")
                             cache.put(path, deps)
-                            cacheOfNulls.remove(path)
-                            file
                         }
                         deps != null -> {
                             // same as before
-                            log.info("[kts] unchanged deps for $path")
-                            null
                         }
                         else -> {
-                            if (cache.remove(path) != null || cacheOfNulls.remove(path)) {
+                            if (cache.remove(path) != null) {
-                                log.info("[kts] removed deps for $path")
+                                log.debug("[kts] removed deps for $path")
-                                file
                             } // cleared
-                            else null // same as before
                         }
                     }
+                    file
                 }
                 else null // not a script
             }
@@ -126,30 +124,15 @@ class KotlinScriptExternalImportsProvider(val project: Project, private val scri
     }
 
     fun invalidateCaches() {
-        cacheLock.write {
+        cacheLock.write(cache::clear)
-            cache.clear()
-            cacheOfNulls.clear()
-        }
-    }
-
-    fun <TF: Any> invalidateCachesFor(vararg files: TF) { invalidateCachesFor(files.asIterable()) }
-
-    fun <TF: Any> invalidateCachesFor(files: Iterable<TF>) {
-        cacheLock.write {
-            files.forEach { file ->
-                val path = getFilePath(file)
-                cache.remove(path)
-                cacheOfNulls.remove(path)
-            }
-        }
     }
 
     fun getKnownCombinedClasspath(): List<File> = cacheLock.read {
-        cache.values.flatMap { it.classpath }
+        cache.values.flatMap { it?.classpath ?: emptyList() }
     }.distinct()
 
     fun getKnownSourceRoots(): List<File> = cacheLock.read {
-        cache.values.flatMap { it.sources }
+        cache.values.flatMap { it?.sources ?: emptyList() }
     }.distinct()
 
     fun <TF: Any> getCombinedClasspathFor(files: Iterable<TF>): List<File> =

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

@@ -50,7 +50,6 @@ import kotlin.concurrent.write
 @Suppress("SimplifyAssertNotNull")
 class KotlinScriptConfigurationManager(
         private val project: Project,
-        private val dumbService: DumbService,
         private val scriptDefinitionProvider: KotlinScriptDefinitionProvider,
         private val scriptExternalImportsProvider: KotlinScriptExternalImportsProvider
 ) {
@@ -60,15 +59,16 @@ class KotlinScriptConfigurationManager(
 
         StartupManager.getInstance(project).runWhenProjectIsInitialized {
             DumbService.getInstance(project).smartInvokeLater {
-                cacheAllScriptsExtraImports()
+                if (cacheAllScriptsExtraImports()) {
-                invalidateLocalCaches()
+                    invalidateLocalCaches()
-                notifyRootsChanged()
+                    notifyRootsChanged()
+                }
             }
         }
 
         project.messageBus.connect().subscribe(VirtualFileManager.VFS_CHANGES, object : BulkFileListener.Adapter() {
             override fun after(events: List<VFileEvent>) {
-                updateExternalImportsCache(events.mapNotNull { it.file }) {
+                if (updateExternalImportsCache(events.mapNotNull { it.file })) {
                     invalidateLocalCaches()
                     notifyRootsChanged()
                 }
@@ -100,9 +100,7 @@ class KotlinScriptConfigurationManager(
                 if (project.isDisposed) return@runWriteAction
 
                 ProjectRootManagerEx.getInstanceEx(project)?.makeRootsChange(EmptyRunnable.getInstance(), false, true)
-                dumbService.runWhenSmart {
+                ScriptDependenciesModificationTracker.getInstance(project).incModificationCount()
-                    ScriptDependenciesModificationTracker.getInstance(project).incModificationCount()
-                }
             }
         }
 
@@ -128,28 +126,21 @@ class KotlinScriptConfigurationManager(
     fun getAllLibrarySourcesScope() = allLibrarySourcesScope.get()
 
     private fun reloadScriptDefinitions() {
-        makeScriptDefsFromTemplatesProviderExtensions(project, { ep, ex -> log.warn("[kts] Error loading definition from ${ep.id}", ex) }).let {
+        val def = makeScriptDefsFromTemplatesProviderExtensions(project, { ep, ex -> log.warn("[kts] Error loading definition from ${ep.id}", ex) })
-            scriptDefinitionProvider.setScriptDefinitions(it)
+        scriptDefinitionProvider.setScriptDefinitions(def)
+    }
+
+    private fun cacheAllScriptsExtraImports(): Boolean = runReadAction {
+        scriptExternalImportsProvider.run {
+            invalidateCaches()
+            cacheExternalImports(
+                    scriptDefinitionProvider.getAllKnownFileTypes()
+                            .flatMap { FileTypeIndex.getFiles(it, GlobalSearchScope.allScope(project)) }
+            ).any()
         }
     }
 
-    private fun cacheAllScriptsExtraImports() {
+    private fun updateExternalImportsCache(files: Iterable<VirtualFile>) = scriptExternalImportsProvider.updateExternalImportsCache(files).any()
-        runReadAction {
-            scriptExternalImportsProvider.apply {
-                invalidateCaches()
-                cacheExternalImports(
-                        scriptDefinitionProvider.getAllKnownFileTypes()
-                                .flatMap { FileTypeIndex.getFiles(it, GlobalSearchScope.allScope(project)) })
-            }
-        }
-    }
-
-    private fun updateExternalImportsCache(files: Iterable<VirtualFile>, onChange: () -> Unit) {
-        val isChanged = scriptExternalImportsProvider.updateExternalImportsCache(files).any()
-        if (isChanged) {
-            onChange()
-        }
-    }
 
     private fun invalidateLocalCaches() {
         allScriptsClasspathCache.clear()