Implement updatable dependencies index with usage API in the environment for using in REPL

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

@@ -20,9 +20,6 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.load.kotlin.JvmVirtualFileFinder
 import org.jetbrains.kotlin.load.kotlin.JvmVirtualFileFinderFactory
 
-class JvmLazyCliVirtualFileFinderFactory(private val makeIndex: () -> JvmDependenciesIndex) : JvmVirtualFileFinderFactory {
+class JvmCliVirtualFileFinderFactory(private val index: JvmDependenciesIndex) : JvmVirtualFileFinderFactory {
-
-    private val index by lazy { makeIndex() }
-
     override fun create(scope: GlobalSearchScope): JvmVirtualFileFinder = JvmCliVirtualFileFinder(index)
 }

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

@@ -21,6 +21,9 @@ import com.intellij.util.containers.IntArrayList
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import java.util.*
+import java.util.concurrent.locks.ReentrantReadWriteLock
+import kotlin.concurrent.read
+import kotlin.concurrent.write
 
 data class JavaRoot(val file: VirtualFile, val type: JavaRoot.RootType, val prefixFqName: FqName? = null) {
     enum class RootType {
@@ -34,10 +37,31 @@ data class JavaRoot(val file: VirtualFile, val type: JavaRoot.RootType, val pref
     }
 }
 
+interface JvmDependenciesIndex {
+
+    val indexedRoots: Sequence<JavaRoot>
+
+    fun <T : Any> findClass(
+            classId: ClassId,
+            acceptedRootTypes: Set<JavaRoot.RootType> = JavaRoot.SourceAndBinary,
+            findClassGivenDirectory: (VirtualFile, JavaRoot.RootType) -> T?
+    ): T?
+
+    fun traverseDirectoriesInPackage(
+            packageFqName: FqName,
+            acceptedRootTypes: Set<JavaRoot.RootType> = JavaRoot.SourceAndBinary,
+            continueSearch: (VirtualFile, JavaRoot.RootType) -> Boolean
+    )
+
+    fun collectKnownClassNamesInPackage(
+            packageFqName: FqName
+    ): Set<String>
+}
+
 // speeds up finding files/classes in classpath/java source roots
 // NOT THREADSAFE, needs to be adapted/removed if we want compiler to be multithreaded
 // the main idea of this class is for each package to store roots which contains it to avoid excessive file system traversal
-class JvmDependenciesIndex(_roots: List<JavaRoot>) {
+class JvmDependenciesIndexImpl(_roots: List<JavaRoot>): JvmDependenciesIndex {
 
     //these fields are computed based on _roots passed to constructor which are filled in later
     private val roots: List<JavaRoot> by lazy { _roots.toList() }
@@ -74,11 +98,12 @@ class JvmDependenciesIndex(_roots: List<JavaRoot>) {
     // helps improve several scenarios, LazyJavaResolverContext.findClassInJava being the most important
     private var lastClassSearch: Pair<FindClassRequest, SearchResult>? = null
 
+    override val indexedRoots by lazy { roots.asSequence() }
 
     // findClassGivenDirectory MUST check whether the class with this classId exists in given package
-    fun <T : Any> findClass(
+    override fun <T : Any> findClass(
             classId: ClassId,
-            acceptedRootTypes: Set<JavaRoot.RootType> = JavaRoot.SourceAndBinary,
+            acceptedRootTypes: Set<JavaRoot.RootType>,
             findClassGivenDirectory: (VirtualFile, JavaRoot.RootType) -> T?
     ): T? {
         return search(FindClassRequest(classId, acceptedRootTypes)) { dir, rootType ->
@@ -87,9 +112,9 @@ class JvmDependenciesIndex(_roots: List<JavaRoot>) {
         }
     }
 
-    fun traverseDirectoriesInPackage(
+    override fun traverseDirectoriesInPackage(
             packageFqName: FqName,
-            acceptedRootTypes: Set<JavaRoot.RootType> = JavaRoot.SourceAndBinary,
+            acceptedRootTypes: Set<JavaRoot.RootType>,
             continueSearch: (VirtualFile, JavaRoot.RootType) -> Boolean
     ) {
         search(TraverseRequest(packageFqName, acceptedRootTypes)) { dir, rootType ->
@@ -97,7 +122,7 @@ class JvmDependenciesIndex(_roots: List<JavaRoot>) {
         }
     }
 
-    fun collectKnownClassNamesInPackage(
+    override fun collectKnownClassNamesInPackage(
             packageFqName: FqName
     ): Set<String> {
         var result = hashSetOf<String>()
@@ -285,5 +310,54 @@ class JvmDependenciesIndex(_roots: List<JavaRoot>) {
     }
 }
 
+class JvmDependenciesDynamicCompoundIndex() : JvmDependenciesIndex {
+
+    private val indices = arrayListOf<JvmDependenciesIndex>()
+    private val lock = ReentrantReadWriteLock()
+
+    fun addIndex(index: JvmDependenciesIndex) {
+        lock.write {
+            indices.add(index)
+        }
+    }
+
+    fun addNewIndexForRoots(roots: Iterable<JavaRoot>): JvmDependenciesIndex? =
+            lock.read {
+                val alreadyIndexed = indexedRoots.toHashSet()
+                val newRoots = roots.filter { !alreadyIndexed.contains(it) }
+                if (newRoots.isEmpty()) null
+                else {
+                    val index = JvmDependenciesIndexImpl(newRoots)
+                    addIndex(index)
+                    index
+                }
+            }
+
+    override val indexedRoots: Sequence<JavaRoot> get() = indices.asSequence().flatMap { it.indexedRoots }
+
+    override fun <T : Any> findClass(
+            classId: ClassId,
+            acceptedRootTypes: Set<JavaRoot.RootType>,
+            findClassGivenDirectory: (VirtualFile, JavaRoot.RootType) -> T?
+    ): T? =
+            lock.read {
+                indices.asSequence().mapNotNull { it.findClass(classId, acceptedRootTypes, findClassGivenDirectory) }.firstOrNull()
+            }
+
+    override fun traverseDirectoriesInPackage(
+            packageFqName: FqName,
+            acceptedRootTypes: Set<JavaRoot.RootType>,
+            continueSearch: (VirtualFile, JavaRoot.RootType) -> Boolean
+    ) {
+        lock.read {
+            indices.forEach { it.traverseDirectoriesInPackage(packageFqName, acceptedRootTypes, continueSearch) }
+        }
+    }
+
+    override fun collectKnownClassNamesInPackage(packageFqName: FqName): Set<String> = lock.read {
+        indices.fold(hashSetOf(), { s, index -> s.addAll(index.collectKnownClassNamesInPackage(packageFqName)); s })
+    }
+}
+
 private fun IntArrayList.lastOrNull() = if (isEmpty) null else get(size() - 1)
 private val IntArrayList.indices: IntRange get() = 0..(size() - 1)

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

@@ -52,6 +52,7 @@ import com.intellij.psi.impl.file.impl.JavaFileManager
 import com.intellij.psi.meta.MetaDataContributor
 import com.intellij.psi.stubs.BinaryFileStubBuilders
 import com.intellij.psi.util.JavaClassSupers
+import com.intellij.util.io.URLUtil
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.asJava.LightClassGenerationSupport
 import org.jetbrains.kotlin.asJava.classes.KtLightClassForFacade
@@ -71,9 +72,7 @@ import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.codegen.extensions.ClassBuilderInterceptorExtension
 import org.jetbrains.kotlin.codegen.extensions.ExpressionCodegenExtension
 import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
-import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.*
-import org.jetbrains.kotlin.config.JVMConfigurationKeys
-import org.jetbrains.kotlin.config.kotlinSourceRoots
 import org.jetbrains.kotlin.extensions.ExternalDeclarationsProvider
 import org.jetbrains.kotlin.extensions.StorageComponentContainerContributor
 import org.jetbrains.kotlin.idea.KotlinFileType
@@ -108,7 +107,7 @@ class KotlinCoreEnvironment private constructor(
         }
     }
     private val sourceFiles = ArrayList<KtFile>()
-    private val javaRoots = ArrayList<JavaRoot>()
+    private val rootsIndex: JvmDependenciesIndex
 
     val configuration: CompilerConfiguration = configuration.copy().apply { setReadOnly(true) }
 
@@ -124,10 +123,6 @@ class KotlinCoreEnvironment private constructor(
         registerProjectServicesForCLI(projectEnvironment)
         registerProjectServices(projectEnvironment)
 
-        val fileManager = ServiceManager.getService(project, CoreJavaFileManager::class.java)
-        val index = JvmDependenciesIndex(javaRoots)
-        (fileManager as KotlinCliJavaFileManagerImpl).initIndex(index)
-
         sourceFiles.addAll(CompileEnvironmentUtil.getKtFiles(project, getSourceRootsCheckingForDuplicates(), this.configuration, {
             message ->
             report(ERROR, message)
@@ -158,7 +153,20 @@ class KotlinCoreEnvironment private constructor(
             }
         }
 
-        project.registerService(JvmVirtualFileFinderFactory::class.java, JvmLazyCliVirtualFileFinderFactory({ fillClasspath(configuration); index } ))
+        val initialRoots = configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS).classpathRoots()
+        val initialIndex = JvmDependenciesIndexImpl(initialRoots)
+        updateClasspathFromRootsIndex(initialIndex)
+
+        // replacing index with updatable one only for REPL mode, so we can add roots to classpath then compiling
+        // subsequent lines in REPL
+        rootsIndex = if (!configuration.getBoolean(CommonConfigurationKeys.REPL_MODE)) initialIndex
+                     else JvmDependenciesDynamicCompoundIndex().apply {
+                         addIndex(initialIndex)
+                     }
+        (ServiceManager.getService(project, CoreJavaFileManager::class.java)
+            as KotlinCliJavaFileManagerImpl).initIndex(rootsIndex)
+
+        project.registerService(JvmVirtualFileFinderFactory::class.java, JvmCliVirtualFileFinderFactory(rootsIndex))
 
         ExternalDeclarationsProvider.registerExtensionPoint(project)
         ExpressionCodegenExtension.registerExtensionPoint(project)
@@ -189,33 +197,49 @@ class KotlinCoreEnvironment private constructor(
                 StringUtil.getLineBreakCount(it.text) + (if (StringUtil.endsWithLineBreak(text)) 0 else 1)
             }
 
-    private fun fillClasspath(configuration: CompilerConfiguration) {
+    private fun Iterable<ContentRoot>.classpathRoots(): List<JavaRoot> =
-        for (root in configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS)) {
+        filterIsInstance(JvmContentRoot::class.java).mapNotNull {
-            val javaRoot = root as? JvmContentRoot ?: continue
+            javaRoot -> contentRootToVirtualFile(javaRoot)?.let { virtualFile ->
-            val virtualFile = contentRootToVirtualFile(javaRoot) ?: continue
 
-            projectEnvironment.addSourcesToClasspath(virtualFile)
+                val prefixPackageFqName = (javaRoot as? JavaSourceRoot)?.packagePrefix?.let {
-
+                    if (isValidJavaFqName(it)) {
-            val prefixPackageFqName = (javaRoot as? JavaSourceRoot)?.packagePrefix?.let {
+                        FqName(it)
-                if (isValidJavaFqName(it)) {
+                    }
-                    FqName(it)
+                    else {
+                        report(WARNING, "Invalid package prefix name is ignored: $it")
+                        null
+                    }
                 }
-                else {
+
-                    report(WARNING, "Invalid package prefix name is ignored: $it")
+                val rootType = when (javaRoot) {
-                    null
+                    is JavaSourceRoot -> JavaRoot.RootType.SOURCE
+                    is JvmClasspathRoot -> JavaRoot.RootType.BINARY
+                    else -> throw IllegalStateException()
                 }
-            }
 
-            val rootType = when (javaRoot) {
+                JavaRoot(virtualFile, rootType, prefixPackageFqName)
-                is JavaSourceRoot -> JavaRoot.RootType.SOURCE
-                is JvmClasspathRoot -> JavaRoot.RootType.BINARY
-                else -> throw IllegalStateException()
             }
+        }
 
-            javaRoots.add(JavaRoot(virtualFile, rootType, prefixPackageFqName))
+    private fun updateClasspathFromRootsIndex(index: JvmDependenciesIndex) {
+        index.indexedRoots.forEach {
+            projectEnvironment.addSourcesToClasspath(it.file)
         }
     }
 
+    @Suppress("unused") // used externally
+    fun tryUpdateClasspath(files: Iterable<File>): List<File>? =
+        if (rootsIndex !is JvmDependenciesDynamicCompoundIndex) {
+            report(WARNING, "Unable to update classpath after initialization, it is only allowed in REPL")
+            null
+        }
+        else {
+            rootsIndex.addNewIndexForRoots(files.map { JvmClasspathRoot(it) }.classpathRoots())?.let {
+                updateClasspathFromRootsIndex(it)
+                it.indexedRoots.mapNotNull { File(it.file.path.substringBefore(URLUtil.JAR_SEPARATOR)) }.toList()
+            } ?: emptyList()
+        }
+
     fun contentRootToVirtualFile(root: JvmContentRoot): VirtualFile? {
         when (root) {
             is JvmClasspathRoot -> {
@@ -240,7 +264,7 @@ class KotlinCoreEnvironment private constructor(
 
     private fun findJarRoot(root: JvmClasspathRoot): VirtualFile? {
         val path = root.file
-        val jarFile = applicationEnvironment.jarFileSystem.findFileByPath("${path}!/")
+        val jarFile = applicationEnvironment.jarFileSystem.findFileByPath("${path}${URLUtil.JAR_SEPARATOR}")
         if (jarFile == null) {
             report(WARNING, "Classpath entry points to a file that is not a JAR archive: $path")
             return null

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/ReplInterpreter.kt

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.codegen.ClassBuilderFactories
 import org.jetbrains.kotlin.codegen.CompilationErrorHandler
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
@@ -64,6 +65,7 @@ class ReplInterpreter(
 
     private val environment = run {
         configuration.add(JVMConfigurationKeys.SCRIPT_DEFINITIONS, REPL_LINE_AS_SCRIPT_DEFINITION)
+        configuration.put(CommonConfigurationKeys.REPL_MODE, true)
         KotlinCoreEnvironment.createForProduction(disposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
     }
 

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

@@ -28,4 +28,7 @@ public class CommonConfigurationKeys {
 
     public static final CompilerConfigurationKey<String> MODULE_NAME =
             CompilerConfigurationKey.create("module name");
+
+    public static final CompilerConfigurationKey<Boolean> REPL_MODE =
+            CompilerConfigurationKey.create("REPL mode");
 }

compiler/tests/org/jetbrains/kotlin/cli/jvm/KotlinCliJavaFileManagerTest.kt

@@ -196,7 +196,7 @@ class KotlinCliJavaFileManagerTest : KotlinTestWithEnvironment() {
         val coreJavaFileManager = ServiceManager.getService(project, CoreJavaFileManager::class.java) as KotlinCliJavaFileManagerImpl
 
         val root = environment.contentRootToVirtualFile(JavaSourceRoot(javaFilesDir!!, null))!!
-        coreJavaFileManager.initIndex(JvmDependenciesIndex(listOf(JavaRoot(root, JavaRoot.RootType.SOURCE))))
+        coreJavaFileManager.initIndex(JvmDependenciesIndexImpl(listOf(JavaRoot(root, JavaRoot.RootType.SOURCE))))
 
         return coreJavaFileManager
     }