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 =
-        ConcurrentFactoryMap.createMap<String, CoreJrtHandler?> { jdkHomePath ->
+    private val roots =
+        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? {
-        val parentPath = path.parent
-        return if (parentPath != null) CoreJrtVirtualFile(handler, parentPath) else null
-    }
+    override fun getParent(): VirtualFile? = parent
 
-    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
-            else -> paths.map { path -> CoreJrtVirtualFile(handler, path) }.toTypedArray()
+            paths.isEmpty() -> EMPTY_ARRAY
+            else -> paths.map { path -> CoreJrtVirtualFile(virtualFileSystem, jdkHomePath, path, parent = this) }.toTypedArray()
         }
     }