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/cli/src/org/jetbrains/kotlin/cli/common/script/CliScriptDefinitionProvider.kt

@@ -16,24 +16,15 @@
 
 package org.jetbrains.kotlin.cli.common.script
 
-import org.jetbrains.kotlin.script.KotlinScriptDefinition
-import org.jetbrains.kotlin.script.ScriptDefinitionProvider
-import org.jetbrains.kotlin.script.StandardScriptDefinition
-import java.util.concurrent.locks.ReentrantReadWriteLock
-import kotlin.concurrent.read
+import org.jetbrains.kotlin.script.*
 import kotlin.concurrent.write
 
-class CliScriptDefinitionProvider : ScriptDefinitionProvider {
+class CliScriptDefinitionProvider : LazyScriptDefinitionProvider() {
+    private val definitionsFromSources: MutableList<Sequence<KotlinScriptDefinition>> = arrayListOf()
     private val definitions: MutableList<KotlinScriptDefinition> = arrayListOf(StandardScriptDefinition)
-    private val lock = ReentrantReadWriteLock()
 
-    override fun findScriptDefinition(fileName: String) = lock.read {
-        definitions.firstOrNull { it.isScript(fileName) }
-    }
-
-    override fun isScript(fileName: String) = lock.read {
-        definitions.any { it.isScript(fileName) }
-    }
+    override val currentDefinitions: Sequence<KotlinScriptDefinition> =
+        definitionsFromSources.asSequence().flatMap { it } + definitions.asSequence()
 
     fun setScriptDefinitions(newDefinitions: List<KotlinScriptDefinition>) {
         lock.write {
@@ -41,4 +32,13 @@ class CliScriptDefinitionProvider : ScriptDefinitionProvider {
             definitions.addAll(newDefinitions)
         }
     }
+
+    fun setScriptDefinitionsSources(newSources: List<ScriptDefinitionsSource>) {
+        lock.write {
+            definitionsFromSources.clear()
+            for (it in newSources) {
+                definitionsFromSources.add(it.definitions.constrainOnce())
+            }
+        }
+    }
 }

compiler/cli/src/org/jetbrains/kotlin/cli/common/script/CliScriptDependenciesProvider.kt

@@ -19,10 +19,8 @@ package org.jetbrains.kotlin.cli.common.script
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
-import org.jetbrains.kotlin.script.ScriptContentLoader
-import org.jetbrains.kotlin.script.ScriptDefinitionProvider
-import org.jetbrains.kotlin.script.ScriptDependenciesProvider
-import org.jetbrains.kotlin.script.adjustByDefinition
+import org.jetbrains.kotlin.script.findScriptDefinition
+import org.jetbrains.kotlin.script.*
 import java.io.File
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.read

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

@@ -209,6 +209,9 @@ class KotlinCoreEnvironment private constructor(
 
         val scriptDefinitionProvider = ScriptDefinitionProvider.getInstance(project) as? CliScriptDefinitionProvider
         if (scriptDefinitionProvider != null) {
+            scriptDefinitionProvider.setScriptDefinitionsSources(
+                configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES)
+            )
             scriptDefinitionProvider.setScriptDefinitions(
                     configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS))
 

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

@@ -210,6 +210,9 @@ class KotlinCoreEnvironment private constructor(
 
         val scriptDefinitionProvider = ScriptDefinitionProvider.getInstance(project) as? CliScriptDefinitionProvider
         if (scriptDefinitionProvider != null) {
+            scriptDefinitionProvider.setScriptDefinitionsSources(
+                configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES)
+            )
             scriptDefinitionProvider.setScriptDefinitions(
                     configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS))
 

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

@@ -212,6 +212,9 @@ class KotlinCoreEnvironment private constructor(
 
         val scriptDefinitionProvider = ScriptDefinitionProvider.getInstance(project) as? CliScriptDefinitionProvider
         if (scriptDefinitionProvider != null) {
+            scriptDefinitionProvider.setScriptDefinitionsSources(
+                configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES)
+            )
             scriptDefinitionProvider.setScriptDefinitions(
                     configuration.getList(JVMConfigurationKeys.SCRIPT_DEFINITIONS))
 

compiler/frontend.java/src/org/jetbrains/kotlin/config/JVMConfigurationKeys.java

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.load.java.JavaClassesTracker;
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents;
 import org.jetbrains.kotlin.modules.Module;
 import org.jetbrains.kotlin.script.KotlinScriptDefinition;
+import org.jetbrains.kotlin.script.ScriptDefinitionsSource;
 
 import java.io.File;
 import java.util.List;
@@ -49,6 +50,9 @@ public class JVMConfigurationKeys {
     public static final CompilerConfigurationKey<List<KotlinScriptDefinition>> SCRIPT_DEFINITIONS =
             CompilerConfigurationKey.create("script definitions");
 
+    public static final CompilerConfigurationKey<List<ScriptDefinitionsSource>> SCRIPT_DEFINITIONS_SOURCES =
+            CompilerConfigurationKey.create("script definitions sources");
+
     public static final CompilerConfigurationKey<Boolean> DISABLE_STANDARD_SCRIPT_DEFINITION =
             CompilerConfigurationKey.create("Disable standard kotlin script support");
 

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>
+}

compiler/tests/org/jetbrains/kotlin/scripts/ScriptProviderTest.kt

@@ -0,0 +1,80 @@
+/*
+ * 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.scripts
+
+import org.jetbrains.kotlin.cli.common.script.CliScriptDefinitionProvider
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.ScriptDefinitionsSource
+import org.jetbrains.kotlin.test.testFramework.KtUsefulTestCase
+import org.junit.Assert
+import org.junit.Test
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.coroutines.experimental.buildSequence
+import kotlin.script.templates.standard.ScriptTemplateWithArgs
+
+class ScriptProviderTest : KtUsefulTestCase() {
+
+    @Test
+    fun testLazyScriptDefinitionsProvider() {
+
+        val genDefCounter = AtomicInteger()
+        val standardDef = FakeScriptDefinition()
+        val shadedDef = FakeScriptDefinition(".x.kts")
+        val provider = CliScriptDefinitionProvider().apply {
+            setScriptDefinitions(listOf(shadedDef, standardDef))
+            setScriptDefinitionsSources(listOf(TestScriptDefinitionSource(genDefCounter, ".y.kts", ".x.kts")))
+        }
+
+        Assert.assertEquals(0, genDefCounter.get())
+
+        provider.isScript("a.kt").let {
+            Assert.assertFalse(it)
+            Assert.assertEquals(0, genDefCounter.get())
+        }
+
+        provider.isScript("a.y.kts").let {
+            Assert.assertTrue(it)
+            Assert.assertEquals(1, genDefCounter.get())
+        }
+
+        provider.isScript("a.x.kts").let {
+            Assert.assertTrue(it)
+            Assert.assertEquals(2, genDefCounter.get())
+            Assert.assertEquals(0, shadedDef.matchCounter.get())
+        }
+
+        provider.isScript("a.z.kts").let {
+            Assert.assertTrue(it)
+            Assert.assertEquals(2, genDefCounter.get())
+            Assert.assertEquals(1, standardDef.matchCounter.get())
+        }
+
+        provider.isScript("a.ktx").let {
+            Assert.assertFalse(it)
+            Assert.assertEquals(2, genDefCounter.get())
+        }
+    }
+}
+
+private class FakeScriptDefinition(val suffix: String = ".kts") : KotlinScriptDefinition(ScriptTemplateWithArgs::class) {
+    val matchCounter = AtomicInteger()
+    override fun isScript(fileName: String): Boolean = fileName.endsWith(suffix).also {
+        if (it) matchCounter.incrementAndGet()
+    }
+}
+
+private class TestScriptDefinitionSource(val counter: AtomicInteger, val defGens: Iterable<() -> FakeScriptDefinition>) :
+    ScriptDefinitionsSource
+{
+    constructor(counter: AtomicInteger, vararg suffixes: String) : this(counter, suffixes.map { { FakeScriptDefinition(it) } })
+
+    override val definitions: Sequence<KotlinScriptDefinition> = buildSequence {
+        for (gen in defGens) {
+            counter.incrementAndGet()
+            yield(gen())
+        }
+    }
+}