Support async script dependency updates for annotation based templates

Do not try to cache dependencies for every script on startup
Schedule cache updates upon request as opposed to updating them synchronously

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

@@ -123,28 +123,35 @@ open class KotlinScriptDefinitionFromAnnotatedTemplate(
     override fun getScriptName(script: KtScript): Name = NameUtils.getScriptNameForFile(script.containingKtFile.name)
 
     override fun <TF: Any> getDependenciesFor(file: TF, project: Project, previousDependencies: KotlinScriptExternalDependencies?): KotlinScriptExternalDependencies? {
-
-        fun makeScriptContents() = BasicScriptContents(file, getAnnotations = {
-            val classLoader = template.java.classLoader
-            takeUnlessError(reportError = false) {
-                getAnnotationEntries(file, project)
-                        .mapNotNull { psiAnn ->
-                            // TODO: consider advanced matching using semantic similar to actual resolving
-                            acceptedAnnotations.find { ann ->
-                                psiAnn.typeName.let { it == ann.simpleName || it == ann.qualifiedName }
-                            }?.let { constructAnnotation(psiAnn, classLoader.loadClass(it.qualifiedName).kotlin as KClass<out Annotation>) }
-                        }
-            }
-            ?: emptyList()
-        })
-
         return takeUnlessError(reportError = false) {
-            val fileDeps = resolver?.resolve(makeScriptContents(), environment, ::logScriptDefMessage, previousDependencies)
+            val fileDeps = resolver?.resolve(makeScriptContents(file, project), environment, ::logScriptDefMessage, previousDependencies)
             // TODO: use it as a Future
             fileDeps?.get()
         }
     }
 
+    fun <TF: Any>  makeScriptContents(file: TF, project: Project) = BasicScriptContents(file, getAnnotations = {
+        val classLoader = template.java.classLoader
+        takeUnlessError(reportError = false) {
+            getAnnotationEntries(file, project)
+                    .mapNotNull { psiAnn ->
+                        // TODO: consider advanced matching using semantic similar to actual resolving
+                        acceptedAnnotations.find { ann ->
+                            psiAnn.typeName.let { it == ann.simpleName || it == ann.qualifiedName }
+                        }?.let { constructAnnotation(psiAnn, classLoader.loadClass(it.qualifiedName).kotlin as KClass<out Annotation>) }
+                    }
+        }
+        ?: emptyList()
+    })
+
+
+    fun getDependenciesFor(previousDependencies: KotlinScriptExternalDependencies?, scriptContents: ScriptContents): KotlinScriptExternalDependencies? {
+        return takeUnlessError(reportError = false) {
+            // TODO: use it as a Future
+            resolver?.resolve(scriptContents, environment, ::logScriptDefMessage, previousDependencies)?.get()
+        }
+    }
+
     private fun <TF: Any> getAnnotationEntries(file: TF, project: Project): Iterable<KtAnnotationEntry> = when (file) {
         is PsiFile -> getAnnotationEntriesFromPsiFile(file)
         is VirtualFile -> getAnnotationEntriesFromVirtualFile(file, project)

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

@@ -21,12 +21,10 @@ import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.Extensions
-import com.intellij.openapi.project.DumbService
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectUtil
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.roots.ex.ProjectRootManagerEx
-import com.intellij.openapi.startup.StartupManager
 import com.intellij.openapi.util.EmptyRunnable
 import com.intellij.openapi.vfs.StandardFileSystems
 import com.intellij.openapi.vfs.VirtualFile
@@ -34,14 +32,13 @@ 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.PsiElementFinder
-import com.intellij.psi.search.FileTypeIndex
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.NonClasspathDirectoriesScope
 import com.intellij.util.io.URLUtil
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.script.*
 import java.io.File
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.read
 import kotlin.concurrent.write
@@ -65,16 +62,10 @@ class KotlinScriptConfigurationManager(
 
     init {
         reloadScriptDefinitions()
+        listenToVfsChanges()
+    }
 
-        StartupManager.getInstance(project).runWhenProjectIsInitialized {
+    private fun listenToVfsChanges() {
-            DumbService.getInstance(project).smartInvokeLater {
-                if (populateCache()) {
-                    invalidateLocalCaches()
-                    notifyRootsChanged()
-                }
-            }
-        }
-
         project.messageBus.connect().subscribe(VirtualFileManager.VFS_CHANGES, object : BulkFileListener.Adapter() {
 
             val projectFileIndex = ProjectRootManager.getInstance(project).fileIndex
@@ -97,7 +88,7 @@ class KotlinScriptConfigurationManager(
     }
 
     private val allScriptsClasspathCache = ClearableLazyValue(cacheLock) {
-        val files = cache.values.flatMap { it?.classpath ?: emptyList() }.distinct()
+        val files = cache.values.flatMap { it.classpath }.distinct()
         toVfsRoots(files)
     }
 
@@ -106,7 +97,7 @@ class KotlinScriptConfigurationManager(
     }
 
     private val allLibrarySourcesCache = ClearableLazyValue(cacheLock) {
-        toVfsRoots(cache.values.flatMap { it?.sources ?: emptyList() }.distinct())
+        toVfsRoots(cache.values.flatMap { it.sources }.distinct())
     }
 
     private val allLibrarySourcesScope = ClearableLazyValue(cacheLock) {
@@ -132,9 +123,7 @@ class KotlinScriptConfigurationManager(
         }
     }
 
-    fun getScriptClasspath(file: VirtualFile): List<VirtualFile> = toVfsRoots(
+    fun getScriptClasspath(file: VirtualFile): List<VirtualFile> = toVfsRoots(getExternalImports(file).classpath)
-            getExternalImports(file)?.classpath ?: emptyList()
-    )
 
     fun getAllScriptsClasspath(): List<VirtualFile> = allScriptsClasspathCache.get()
 
@@ -149,48 +138,20 @@ class KotlinScriptConfigurationManager(
         scriptDefinitionProvider.setScriptDefinitions(def)
     }
 
-    private fun populateCache(): Boolean {
+    private val cache = hashMapOf<String, KotlinScriptExternalDependencies>()
-        cacheLock.write(cache::clear)
+    // TODO: this is very primitive way of synchronizing these updates
-        return cacheExternalImports(
+    private var lastStamp: TimeStamp = TimeStamps.next()
-                scriptDefinitionProvider.getAllKnownFileTypes()
-                        .flatMap { FileTypeIndex.getFiles(it, GlobalSearchScope.allScope(project)) }
-        ).any()
-    }
 
-    private val cache = hashMapOf<String, KotlinScriptExternalDependencies?>()
+    fun <TF : Any> getExternalImports(file: TF): KotlinScriptExternalDependencies = cacheLock.read {
-
-    fun <TF : Any> getExternalImports(file: TF): KotlinScriptExternalDependencies? = cacheLock.read {
         val path = getFilePath(file)
         cache[path]?.let { return it }
 
-        val scriptDef = scriptDefinitionProvider.findScriptDefinition(file) ?: return null
+        updateExternalImportsCache(listOf(file))
-        return scriptDef.getDependenciesFor(file, project, null).also {
+
-            cacheLock.write {
+        return cache[path] ?: NoDependencies
-                cache.put(path, it)
-            }
-        }
     }
 
-    // optimized for initial caching, additional handling of possible duplicates to save a call to distinct
+    private fun <TF: Any> updateExternalImportsCache(files: Iterable<TF>) = cacheLock.write {
-    // returns list of cached files
-    fun <TF : Any> cacheExternalImports(files: Iterable<TF>): Iterable<TF> = cacheLock.write {
-        return files.mapNotNull { file ->
-            scriptDefinitionProvider.findScriptDefinition(file)?.let {
-                cacheForFile(file, getFilePath(file), it)
-            }
-        }
-    }
-
-    private fun <TF : Any> cacheForFile(file: TF, path: String, scriptDef: KotlinScriptDefinition): TF? {
-        if (!isValidFile(file) || cache.containsKey(path)) return null
-
-        val deps = scriptDef.getDependenciesFor(file, project, null)
-        cache.put(path, deps)
-
-        return file.takeIf { deps != null }
-    }
-
-    private fun updateExternalImportsCache(files: Iterable<VirtualFile>) = cacheLock.write {
         files.mapNotNull { file ->
             scriptDefinitionProvider.findScriptDefinition(file)?.let {
                 updateForFile(file, it)
@@ -203,9 +164,35 @@ class KotlinScriptConfigurationManager(
             return cache.remove(getFilePath(file)) != null
         }
 
+        // TODO: support apis that allow for async updates for any template
+        if (scriptDef is KotlinScriptDefinitionFromAnnotatedTemplate) {
+            return updateAsync(file, scriptDef)
+        }
         return updateSync(file, scriptDef)
     }
 
+    fun <TF : Any> updateAsync(file: TF, scriptDefinition: KotlinScriptDefinitionFromAnnotatedTemplate): Boolean {
+        val path = getFilePath(file)
+        val oldDependencies = cache[path]
+
+        val scriptContents = scriptDefinition.makeScriptContents(file, project)
+        val updateStamp = TimeStamps.next()
+
+        CompletableFuture.supplyAsync {
+            val newDependencies = scriptDefinition.getDependenciesFor(oldDependencies, scriptContents)
+            cacheLock.write {
+                if (updateStamp >= lastStamp) {
+                    lastStamp = updateStamp
+                    if (cacheNewDependencies(newDependencies, oldDependencies, path)) {
+                        invalidateLocalCaches()
+                        notifyRootsChanged()
+                    }
+                }
+            }
+        }
+        return false // not changed immediately
+    }
+
     private fun <TF : Any> updateSync(file: TF, scriptDef: KotlinScriptDefinition): Boolean {
         val path = getFilePath(file)
         val oldDeps = cache[path]
@@ -275,6 +262,8 @@ class KotlinScriptConfigurationManager(
             }
         }
     }
+
+    private object NoDependencies: KotlinScriptExternalDependencies
 }
 
 private class ClearableLazyValue<out T : Any>(private val lock: ReentrantReadWriteLock, private val compute: () -> T) {
@@ -309,3 +298,13 @@ private fun Iterable<File>.isSamePathListAs(other: Iterable<File>): Boolean =
             }
             !(first.hasNext() || second.hasNext())
         }
+
+data class TimeStamp(private val stamp: Long) {
+    operator fun compareTo(other: TimeStamp) = this.stamp.compareTo(other.stamp)
+}
+
+object TimeStamps {
+    private var current: Long = 0
+
+    fun next() = TimeStamp(current++)
+}