Optimize CoreJrtFileSystem

Do not recompute children on each file request
At the same time, we should clear the roots after each compilation
just the same way as we do for common jars

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/modules/CoreJrtFileSystem.kt

@@ -26,13 +26,11 @@ import java.io.File
 import java.net.URI
 import java.net.URLClassLoader
 import java.nio.file.FileSystems
-import java.nio.file.Files
-import java.nio.file.Path
 
 // There's JrtFileSystem in idea-full which we can't use in the compiler because it depends on NewVirtualFileSystem, absent in intellij-core
 class CoreJrtFileSystem : DeprecatedVirtualFileSystem() {
-    private val handlers =
+    private val roots =
-        ConcurrentFactoryMap.createMap<String, CoreJrtHandler?> { jdkHomePath ->
+        ConcurrentFactoryMap.createMap<String, CoreJrtVirtualFile?> { jdkHomePath ->
             val jdkHome = File(jdkHomePath)
             val rootUri = URI.create(StandardFileSystems.JRT_PROTOCOL + ":/")
             val jrtFsJar = loadJrtFsJar(jdkHome) ?: return@createMap null
@@ -50,31 +48,28 @@ class CoreJrtFileSystem : DeprecatedVirtualFileSystem() {
                     }
                     FileSystems.newFileSystem(rootUri, emptyMap<String, Nothing>(), classLoader)
                 }
-            CoreJrtHandler(this, jdkHomePath, fileSystem.getPath(""))
+            CoreJrtVirtualFile(this, jdkHomePath, fileSystem.getPath(""), parent = null)
         }
 
-    internal class CoreJrtHandler(
-        val virtualFileSystem: CoreJrtFileSystem,
-        val jdkHomePath: String,
-        private val root: Path
-    ) {
-        fun findFile(fileName: String): VirtualFile? {
-            val path = root.resolve(fileName)
-            return if (Files.exists(path)) CoreJrtVirtualFile(this, path) else null
-        }
-    }
-
     override fun getProtocol(): String = StandardFileSystems.JRT_PROTOCOL
 
     override fun findFileByPath(path: String): VirtualFile? {
         val (jdkHomePath, pathInImage) = splitPath(path)
-        return handlers[jdkHomePath]?.findFile(pathInImage)
+        val root = roots[jdkHomePath] ?: return null
+
+        if (pathInImage.isEmpty()) return root
+
+        return root.findFileByRelativePath(pathInImage)
     }
 
     override fun refresh(asynchronous: Boolean) {}
 
     override fun refreshAndFindFileByPath(path: String): VirtualFile? = findFileByPath(path)
 
+    fun clearRoots() {
+        roots.clear()
+    }
+
     companion object {
         private fun loadJrtFsJar(jdkHome: File): File? =
             File(jdkHome, "lib/jrt-fs.jar").takeIf(File::exists)

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/modules/CoreJrtVirtualFile.kt

@@ -21,7 +21,6 @@ import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileSystem
 import com.intellij.util.io.URLUtil
-import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem.CoreJrtHandler
 import java.io.IOException
 import java.io.InputStream
 import java.io.OutputStream
@@ -29,17 +28,22 @@ import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.attribute.BasicFileAttributes
 
-internal class CoreJrtVirtualFile(private val handler: CoreJrtHandler, private val path: Path) : VirtualFile() {
+internal class CoreJrtVirtualFile(
+    private val virtualFileSystem: CoreJrtFileSystem,
+    private val jdkHomePath: String,
+    private val path: Path,
+    private val parent: CoreJrtVirtualFile?,
+) : VirtualFile() {
     // TODO: catch IOException?
     private val attributes: BasicFileAttributes get() = Files.readAttributes(path, BasicFileAttributes::class.java)
 
-    override fun getFileSystem(): VirtualFileSystem = handler.virtualFileSystem
+    override fun getFileSystem(): VirtualFileSystem = virtualFileSystem
 
     override fun getName(): String =
         path.fileName.toString()
 
     override fun getPath(): String =
-        FileUtil.toSystemIndependentName(handler.jdkHomePath + URLUtil.JAR_SEPARATOR + path)
+        FileUtil.toSystemIndependentName(jdkHomePath + URLUtil.JAR_SEPARATOR + path)
 
     override fun isWritable(): Boolean = false
 
@@ -47,20 +51,21 @@ internal class CoreJrtVirtualFile(private val handler: CoreJrtHandler, private v
 
     override fun isValid(): Boolean = true
 
-    override fun getParent(): VirtualFile? {
+    override fun getParent(): VirtualFile? = parent
-        val parentPath = path.parent
-        return if (parentPath != null) CoreJrtVirtualFile(handler, parentPath) else null
-    }
 
-    override fun getChildren(): Array<out VirtualFile> {
+    private val myChildren by lazy { computeChildren() }
+
+    override fun getChildren(): Array<out VirtualFile> = myChildren
+
+    private fun computeChildren(): Array<out VirtualFile> {
         val paths = try {
             Files.newDirectoryStream(path).use(Iterable<Path>::toList)
         } catch (e: IOException) {
             emptyList<Path>()
         }
         return when {
-            paths.isEmpty() -> VirtualFile.EMPTY_ARRAY
+            paths.isEmpty() -> EMPTY_ARRAY
-            else -> paths.map { path -> CoreJrtVirtualFile(handler, path) }.toTypedArray()
+            else -> paths.map { path -> CoreJrtVirtualFile(virtualFileSystem, jdkHomePath, path, parent = this) }.toTypedArray()
         }
     }
 

compiler/daemon/src/org/jetbrains/kotlin/daemon/CompileServiceImpl.kt

@@ -32,17 +32,16 @@ import org.jetbrains.kotlin.cli.common.messages.MessageCollector
 import org.jetbrains.kotlin.cli.common.repl.ReplCheckResult
 import org.jetbrains.kotlin.cli.common.repl.ReplCodeLine
 import org.jetbrains.kotlin.cli.common.repl.ReplCompileResult
-import org.jetbrains.kotlin.cli.common.repl.ReplEvalResult
 import org.jetbrains.kotlin.cli.js.K2JSCompiler
 import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 import org.jetbrains.kotlin.cli.metadata.K2MetadataCompiler
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import org.jetbrains.kotlin.config.Services
 import org.jetbrains.kotlin.daemon.common.*
 import org.jetbrains.kotlin.daemon.report.CompileServicesFacadeMessageCollector
 import org.jetbrains.kotlin.daemon.report.DaemonMessageReporter
-import org.jetbrains.kotlin.daemon.report.DaemonMessageReporterPrintStreamAdapter
 import org.jetbrains.kotlin.daemon.report.getBuildReporter
 import org.jetbrains.kotlin.incremental.*
 import org.jetbrains.kotlin.incremental.components.EnumWhenTracker
@@ -57,9 +56,7 @@ import org.jetbrains.kotlin.incremental.multiproject.ModulesApiHistoryJvm
 import org.jetbrains.kotlin.incremental.parsing.classesFqNames
 import org.jetbrains.kotlin.load.kotlin.incremental.components.IncrementalCompilationComponents
 import org.jetbrains.kotlin.progress.CompilationCanceledStatus
-import java.io.BufferedOutputStream
 import java.io.File
-import java.io.PrintStream
 import java.rmi.NoSuchObjectException
 import java.rmi.registry.Registry
 import java.rmi.server.UnicastRemoteObject
@@ -1073,6 +1070,7 @@ class CompileServiceImpl(
         ZipHandler.clearFileAccessorCache()
         KotlinCoreEnvironment.applicationEnvironment?.apply {
             (jarFileSystem as? CoreJarFileSystem)?.clearHandlersCache()
+            (jrtFileSystem as? CoreJrtFileSystem)?.clearRoots()
             idleCleanup()
         }
     }