Compute module mappings eagerly in JvmPackagePartProvider, refactor

Previously we traversed all notLoadedRoots on each request for package
parts with the given package FQ name. Since notLoadedRoots might contain
a lot of roots (which never transition into "loadedModules" because e.g.
they are not Kotlin libraries, but just Java libraries or SDK roots with
the META-INF directory), this was potentially hurting performance. It
seems it's more optimal to compute everything eagerly once
JvmPackagePartProvider is constructed.

Another problem with the previous version of JvmPackagePartProvider was
that it did not support "updateable classpath" which is used by REPL and
kapt2, it only used the initial roots provided in the
CompilerConfiguration. In REPL specifically, we would thus fail to
resolve top-level callables from libraries which were dynamically added
to the execution classpath (via some kind of a @DependsOn annotation).
In the new code, JvmPackagePartProvider no longer depends on
CompilerConfiguration to avoid this sort of confusion, but rather relies
on the object that constructed it (KotlinCoreEnvironment in this case)
to provide the correct roots. This is also beneficial because the
computation of actual VirtualFile-based roots from the ones in the
CompilerConfiguration might get trickier with modular Java 9 roots

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

@@ -19,9 +19,8 @@ package org.jetbrains.kotlin.cli.jvm.compiler
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.util.SmartList
-import org.jetbrains.kotlin.cli.jvm.config.JvmClasspathRoot
+import org.jetbrains.kotlin.cli.jvm.index.JavaRoot
-import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.PackagePartProvider
 import org.jetbrains.kotlin.load.kotlin.ModuleMapping
 import org.jetbrains.kotlin.load.kotlin.PackageParts
@@ -29,20 +28,12 @@ import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
 import java.io.EOFException
 
 class JvmPackagePartProvider(
-        private val env: KotlinCoreEnvironment,
+        languageVersionSettings: LanguageVersionSettings,
         private val scope: GlobalSearchScope
 ) : PackagePartProvider {
     private data class ModuleMappingInfo(val root: VirtualFile, val mapping: ModuleMapping)
 
-    private val deserializationConfiguration = CompilerDeserializationConfiguration(env.configuration.languageVersionSettings)
+    private val deserializationConfiguration = CompilerDeserializationConfiguration(languageVersionSettings)
-
-    private val notLoadedRoots by lazy(LazyThreadSafetyMode.NONE) {
-        env.configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS)
-                .filterIsInstance<JvmClasspathRoot>()
-                .mapNotNull { env.contentRootToVirtualFile(it) }
-                .filter { it in scope && it.findChild("META-INF") != null }
-                .toMutableList()
-    }
 
     private val loadedModules: MutableList<ModuleMappingInfo> = SmartList()
 
@@ -69,8 +60,6 @@ class JvmPackagePartProvider(
 
     @Synchronized
     private fun getPackageParts(packageFqName: String): Map<VirtualFile, PackageParts> {
-        processNotLoadedRelevantRoots(packageFqName)
-
         val result = mutableMapOf<VirtualFile, PackageParts>()
         for ((root, mapping) in loadedModules) {
             val newParts = mapping.findPackageParts(packageFqName) ?: continue
@@ -79,32 +68,20 @@ class JvmPackagePartProvider(
         return result
     }
 
-    private fun processNotLoadedRelevantRoots(packageFqName: String) {
+    fun addRoots(roots: List<JavaRoot>) {
-        if (notLoadedRoots.isEmpty()) return
+        for ((root, type) in roots) {
+            if (type != JavaRoot.RootType.BINARY) continue
+            if (root !in scope) continue
 
-        val pathParts = packageFqName.split('.')
-
-        val relevantRoots = notLoadedRoots.filter {
-            //filter all roots by package path existing
-            pathParts.fold(it) {
-                parent, part ->
-                if (part.isEmpty()) parent
-                else parent.findChild(part) ?: return@filter false
-            }
-            true
-        }
-        notLoadedRoots.removeAll(relevantRoots)
-
-        for (root in relevantRoots) {
             val metaInf = root.findChild("META-INF") ?: continue
-            val moduleFiles = metaInf.children.filter { it.name.endsWith(ModuleMapping.MAPPING_FILE_EXT) }
+            for (moduleFile in metaInf.children) {
-            for (moduleFile in moduleFiles) {
+                if (!moduleFile.name.endsWith(ModuleMapping.MAPPING_FILE_EXT)) continue
+
                 val mapping = try {
                     ModuleMapping.create(moduleFile.contentsToByteArray(), moduleFile.toString(), deserializationConfiguration)
                 }
                 catch (e: EOFException) {
-                    throw RuntimeException("Error on reading package parts for '$packageFqName' package in '$moduleFile', " +
+                    throw RuntimeException("Error on reading package parts from $moduleFile in $root", e)
-                                           "roots: $notLoadedRoots", e)
                 }
                 loadedModules.add(ModuleMappingInfo(root, mapping))
             }

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

@@ -60,6 +60,7 @@ import com.intellij.psi.impl.PsiTreeChangePreprocessor
 import com.intellij.psi.impl.compiled.ClsCustomNavigationPolicy
 import com.intellij.psi.impl.file.impl.JavaFileManager
 import com.intellij.psi.meta.MetaDataContributor
+import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.stubs.BinaryFileStubBuilders
 import com.intellij.psi.util.JavaClassSupers
 import com.intellij.util.io.URLUtil
@@ -142,8 +143,10 @@ class KotlinCoreEnvironment private constructor(
     }
     private val sourceFiles = mutableListOf<KtFile>()
     private val rootsIndex: JvmDependenciesDynamicCompoundIndex
+    private val packagePartProviders = mutableListOf<JvmPackagePartProvider>()
 
     private val classpathRootsResolver: ClasspathRootsResolver
+    private val initialRoots: List<JavaRoot>
 
     val configuration: CompilerConfiguration = configuration.copy()
 
@@ -194,7 +197,7 @@ class KotlinCoreEnvironment private constructor(
 
         classpathRootsResolver = ClasspathRootsResolver(messageCollector, this::contentRootToVirtualFile)
 
-        val initialRoots = classpathRootsResolver.convertClasspathRoots(configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS))
+        initialRoots = classpathRootsResolver.convertClasspathRoots(configuration.getList(JVMConfigurationKeys.CONTENT_ROOTS))
 
         if (!configuration.getBoolean(JVMConfigurationKeys.SKIP_RUNTIME_VERSION_CHECK) && messageCollector != null) {
             JvmRuntimeVersionsConsistencyChecker.checkCompilerClasspathConsistency(
@@ -223,6 +226,13 @@ class KotlinCoreEnvironment private constructor(
         project.registerService(VirtualFileFinderFactory::class.java, finderFactory)
     }
 
+    fun createPackagePartProvider(scope: GlobalSearchScope): JvmPackagePartProvider {
+        return JvmPackagePartProvider(configuration.languageVersionSettings, scope).apply {
+            addRoots(initialRoots)
+            packagePartProviders += this
+        }
+    }
+
     private val VirtualFile.javaFiles: List<VirtualFile>
         get() = mutableListOf<VirtualFile>().apply {
             VfsUtilCore.processFilesRecursively(this@javaFiles) { file ->
@@ -278,8 +288,14 @@ class KotlinCoreEnvironment private constructor(
         project.putUserData(APPEND_JAVA_SOURCE_ROOTS_HANDLER_KEY, appendJavaSourceRootsHandler)
     }
 
-    fun updateClasspath(roots: List<ContentRoot>): List<File>? {
+    fun updateClasspath(contentRoots: List<ContentRoot>): List<File>? {
-        return rootsIndex.addNewIndexForRoots(classpathRootsResolver.convertClasspathRoots(roots))?.let { newIndex ->
+        val newRoots = classpathRootsResolver.convertClasspathRoots(contentRoots)
+
+        for (packagePartProvider in packagePartProviders) {
+            packagePartProvider.addRoots(newRoots)
+        }
+
+        return rootsIndex.addNewIndexForRoots(newRoots)?.let { newIndex ->
             updateClasspathFromRootsIndex(newIndex)
             newIndex.indexedRoots.mapNotNull { (file) -> File(file.path.substringBefore(URLUtil.JAR_SEPARATOR)) }.toList()
         }.orEmpty()

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

@@ -380,7 +380,7 @@ object KotlinToJVMBytecodeCompiler {
                         environment.getSourceFiles(),
                         CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace(),
                         environment.configuration,
-                        { scope -> JvmPackagePartProvider(environment, scope) },
+                        environment::createPackagePartProvider,
                         sourceModuleSearchScope = scope
                 )
             }

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

@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.cli.common.repl.ILineId
 import org.jetbrains.kotlin.cli.common.repl.ReplCodeLine
 import org.jetbrains.kotlin.cli.common.repl.ReplHistory
 import org.jetbrains.kotlin.cli.jvm.compiler.CliLightClassGenerationSupport
-import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.container.get
@@ -65,8 +64,8 @@ class ReplCodeAnalyzer(environment: KotlinCoreEnvironment) {
                 emptyList(),
                 trace,
                 environment.configuration,
-                { scope -> JvmPackagePartProvider(environment, scope) },
+                environment::createPackagePartProvider,
-                { storageManager, files -> ScriptMutableDeclarationProviderFactory() }
+                { _, _ -> ScriptMutableDeclarationProviderFactory() }
         )
 
         this.module = container.get<ModuleDescriptorImpl>()

compiler/cli/src/org/jetbrains/kotlin/cli/metadata/MetadataSerializer.kt

@@ -22,7 +22,6 @@ import org.jetbrains.kotlin.builtins.BuiltInsBinaryVersion
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
-import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.codegen.JvmCodegenUtil
 import org.jetbrains.kotlin.codegen.serializeToByteArray
@@ -66,7 +65,7 @@ open class MetadataSerializer(private val dependOnOldBuiltIns: Boolean) {
         val analyzer = AnalyzerWithCompilerReport(messageCollector)
         analyzer.analyzeAndReport(files, object : AnalyzerWithCompilerReport.Analyzer {
             override fun analyze(): AnalysisResult = DefaultAnalyzerFacade.analyzeFiles(files, moduleName, dependOnOldBuiltIns) {
-                _, content -> JvmPackagePartProvider(environment, content.moduleContentScope)
+                _, content -> environment.createPackagePartProvider(content.moduleContentScope)
             }
         })
 

compiler/tests-common/org/jetbrains/kotlin/checkers/AbstractDiagnosticsTest.kt

@@ -22,7 +22,6 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.analyzer.common.DefaultAnalyzerFacade
 import org.jetbrains.kotlin.cli.jvm.compiler.CliLightClassGenerationSupport
-import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.config.JvmTarget
 import org.jetbrains.kotlin.config.LanguageVersionSettings
@@ -281,7 +280,7 @@ abstract class AbstractDiagnosticsTest : BaseDiagnosticsTest() {
                     files,
                     moduleTrace,
                     environment.configuration.copy().apply { this.languageVersionSettings = languageVersionSettings },
-                    { scope -> JvmPackagePartProvider(environment, scope) }
+                    environment::createPackagePartProvider
             )
         }
 
@@ -313,7 +312,7 @@ abstract class AbstractDiagnosticsTest : BaseDiagnosticsTest() {
                 FileBasedDeclarationProviderFactory(moduleContext.storageManager, files),
                 moduleContentScope,
                 LookupTracker.DO_NOTHING,
-                JvmPackagePartProvider(environment, moduleContentScope),
+                environment.createPackagePartProvider(moduleContentScope),
                 moduleClassResolver,
                 JvmTarget.JVM_1_6,
                 languageVersionSettings

compiler/tests-common/org/jetbrains/kotlin/codegen/GenerationUtils.kt

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.codegen
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.backend.jvm.JvmIrCodegenFactory
 import org.jetbrains.kotlin.cli.common.output.outputUtils.writeAllTo
-import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -48,11 +47,7 @@ object GenerationUtils {
             environment: KotlinCoreEnvironment,
             classBuilderFactory: ClassBuilderFactory = ClassBuilderFactories.TEST
     ): GenerationState =
-            compileFiles(files, environment.configuration, classBuilderFactory) { scope ->
+            compileFiles(files, environment.configuration, classBuilderFactory, environment::createPackagePartProvider)
-                JvmPackagePartProvider(environment, scope)
-            }
-
-
 
     @JvmStatic
     fun compileFiles(

compiler/tests-common/org/jetbrains/kotlin/jvm/compiler/AbstractLoadJavaTest.java

@@ -21,7 +21,10 @@ import com.intellij.openapi.util.io.FileUtil;
 import junit.framework.ComparisonFailure;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.analyzer.AnalysisResult;
-import org.jetbrains.kotlin.cli.jvm.compiler.*;
+import org.jetbrains.kotlin.cli.jvm.compiler.CliLightClassGenerationSupport;
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles;
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment;
+import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM;
 import org.jetbrains.kotlin.cli.jvm.config.JvmContentRootsKt;
 import org.jetbrains.kotlin.config.CompilerConfiguration;
 import org.jetbrains.kotlin.config.ContentRootsKt;
@@ -159,7 +162,7 @@ public abstract class AbstractLoadJavaTest extends TestCaseWithTmpdir {
 
         AnalysisResult result = TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
                 environment.getProject(), environment.getSourceFiles(), new CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace(),
-                configuration, scope -> new JvmPackagePartProvider(environment, scope)
+                configuration, environment::createPackagePartProvider
         );
 
         PackageViewDescriptor packageView = result.getModuleDescriptor().getPackage(TEST_PACKAGE_FQNAME);

compiler/tests-common/org/jetbrains/kotlin/resolve/lazy/JvmResolveUtil.kt

@@ -20,7 +20,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.AnalysisResult
 import org.jetbrains.kotlin.cli.jvm.compiler.CliLightClassGenerationSupport
-import org.jetbrains.kotlin.cli.jvm.compiler.JvmPackagePartProvider
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.config.CompilerConfiguration
@@ -45,9 +44,7 @@ object JvmResolveUtil {
 
     @JvmStatic
     fun analyzeAndCheckForErrors(files: Collection<KtFile>, environment: KotlinCoreEnvironment): AnalysisResult =
-            analyzeAndCheckForErrors(environment.project, files, environment.configuration) { scope ->
+            analyzeAndCheckForErrors(environment.project, files, environment.configuration, environment::createPackagePartProvider)
-                JvmPackagePartProvider(environment, scope)
-            }
 
     @JvmStatic
     fun analyzeAndCheckForErrors(
@@ -79,9 +76,7 @@ object JvmResolveUtil {
 
     @JvmStatic
     fun analyze(files: Collection<KtFile>, environment: KotlinCoreEnvironment, configuration: CompilerConfiguration): AnalysisResult =
-            analyze(environment.project, files, configuration) { scope ->
+            analyze(environment.project, files, configuration, environment::createPackagePartProvider)
-                JvmPackagePartProvider(environment, scope)
-            }
 
     private fun analyze(
             project: Project,

plugins/uast-kotlin/tests/AbstractKotlinUastTest.kt

@@ -11,7 +11,10 @@ import com.intellij.util.io.URLUtil
 import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
 import org.jetbrains.kotlin.cli.common.messages.PrintingMessageCollector
-import org.jetbrains.kotlin.cli.jvm.compiler.*
+import org.jetbrains.kotlin.cli.jvm.compiler.CliLightClassGenerationSupport
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.TopDownAnalyzerFacadeForJVM
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.addKotlinSourceRoot
@@ -49,14 +52,9 @@ abstract class AbstractKotlinUastTest : AbstractUastTest() {
 
         val trace = CliLightClassGenerationSupport.NoScopeRecordCliBindingTrace()
 
-        val kotlinCoreEnvironment = kotlinCoreEnvironment!!
+        val environment = kotlinCoreEnvironment!!
-
         TopDownAnalyzerFacadeForJVM.analyzeFilesWithJavaIntegration(
-                project,
+                project, environment.getSourceFiles(), trace, compilerConfiguration, environment::createPackagePartProvider
-                kotlinCoreEnvironment.getSourceFiles(),
-                trace,
-                compilerConfiguration,
-                { scope -> JvmPackagePartProvider(kotlinCoreEnvironment, scope) }
         )
 
         val vfs = VirtualFileManager.getInstance().getFileSystem(URLUtil.FILE_PROTOCOL)