Implement quick non-script detection and lazy script discovery..

so ".kt" and ".java" files are not considered as scripts and quickly
filtered out, and for the other files the the checks are implemented
using sequences, mechanisms provided to supply script definitions
lazily, and script discovery is implemented using this mechanisms.

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

@@ -20,11 +20,13 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.PsiFile
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.read
+import kotlin.concurrent.write
 
 interface ScriptDefinitionProvider {
     fun findScriptDefinition(fileName: String): KotlinScriptDefinition?
     fun isScript(fileName: String): Boolean
-    fun findScriptDefinition(file: VirtualFile): KotlinScriptDefinition? = findScriptDefinition(file.name)
 
     companion object {
         fun getInstance(project: Project): ScriptDefinitionProvider =
@@ -32,8 +34,79 @@ interface ScriptDefinitionProvider {
     }
 }
 
+fun ScriptDefinitionProvider.findScriptDefinition(file: VirtualFile): KotlinScriptDefinition? = findScriptDefinition(file.name)
+
 fun getScriptDefinition(file: VirtualFile, project: Project): KotlinScriptDefinition? =
     ScriptDefinitionProvider.getInstance(project).findScriptDefinition(file)
 
 fun getScriptDefinition(psiFile: PsiFile): KotlinScriptDefinition? =
     ScriptDefinitionProvider.getInstance(psiFile.project).findScriptDefinition(psiFile.name)
+
+abstract class LazyScriptDefinitionProvider : ScriptDefinitionProvider {
+
+    protected val lock = ReentrantReadWriteLock()
+
+    protected abstract val currentDefinitions: Sequence<KotlinScriptDefinition>
+
+    private var _cachedDefinitions: Sequence<KotlinScriptDefinition>? = null
+    private val cachedDefinitions: Sequence<KotlinScriptDefinition>
+        get() {
+            // assuming it is always called under read lock
+            assert(lock.readLockCount > 0)
+            if (_cachedDefinitions == null) lock.write {
+                _cachedDefinitions = CashingSequence(currentDefinitions.constrainOnce())
+            }
+            return _cachedDefinitions!!
+        }
+
+    protected fun clearCache() {
+        lock.write {
+            _cachedDefinitions = null
+        }
+    }
+
+    protected open fun nonScriptFileName(fileName: String) = nonScriptFilenameSuffixes.any {
+        fileName.endsWith( it, ignoreCase = true)
+    }
+
+    override fun findScriptDefinition(fileName: String): KotlinScriptDefinition? =
+        if (nonScriptFileName(fileName)) null
+        else lock.read {
+            cachedDefinitions.firstOrNull { it.isScript(fileName) }
+        }
+
+    override fun isScript(fileName: String) =
+        if (nonScriptFileName(fileName)) false
+        else lock.read {
+            cachedDefinitions.any { it.isScript(fileName) }
+        }
+
+    companion object {
+        // TODO: find a common place for storing kotlin-related extensions and reuse values from it everywhere
+        protected val nonScriptFilenameSuffixes = arrayOf(".kt", ".java")
+    }
+}
+
+private class CashingSequence<T>(from: Sequence<T>) : Sequence<T> {
+
+    private val lock = ReentrantReadWriteLock()
+    private val sequenceIterator = from.iterator()
+    private val cache = arrayListOf<T>()
+
+    private inner class CashingIterator : Iterator<T> {
+
+        private val cacheIterator: Iterator<T> = cache.iterator()
+        private var cacheRunOut = !cacheIterator.hasNext()
+
+        private fun cacheHasNext() = !cacheRunOut && (cacheIterator.hasNext().also { if (!it) cacheRunOut = true })
+
+        override fun hasNext(): Boolean = lock.read { cacheHasNext() || sequenceIterator.hasNext() }
+
+        override fun next(): T = lock.read {
+            if (cacheHasNext()) cacheIterator.next()
+            else sequenceIterator.next().also { lock.write { cache.add(it) } }
+        }
+    }
+
+    override fun iterator(): Iterator<T> = CashingIterator()
+}

compiler/psi/src/org/jetbrains/kotlin/script/ScriptDefinitionsSource.kt

@@ -0,0 +1,11 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.script
+
+interface ScriptDefinitionsSource {
+
+    val definitions: Sequence<KotlinScriptDefinition>
+}