Avoid PsiFile searching when resolving imported scripts descriptors.

plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/resolve/LazyScriptDescriptor.kt

@@ -47,6 +47,7 @@ import java.io.File
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.FileBasedScriptSource
 import kotlin.script.experimental.host.GetScriptingClass
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.host.getScriptingClass
@@ -174,29 +175,40 @@ class LazyScriptDescriptor(
 
     override fun computeSupertypes() = listOf(baseClassDescriptor()?.defaultType ?: builtIns.anyType)
 
-    // TODO: consider passing ScriptSource to avoid psi file fsearching
     private inner class ImportedScriptDescriptorsFinder {
 
-        val fileManager = VirtualFileManager.getInstance()
+        val localFS by lazy {
-        val localFS = fileManager.getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+            val fileManager = VirtualFileManager.getInstance()
-        val psiManager = PsiManager.getInstance(scriptInfo.script.project)
+            fileManager.getFileSystem(StandardFileSystems.FILE_PROTOCOL) 
+        }
+        val psiManager by lazy { PsiManager.getInstance(scriptInfo.script.project) }
 
-        operator fun invoke(importedScriptFile: File): ScriptDescriptor? {
+        operator fun invoke(importedScript: SourceCode): ScriptDescriptor? {
+            // Note: is not an error now - if import references other valid source file, it is simply compiled along with script
+            // TODO: check if this is the behavior we want to have - see #KT-28916
+            val ktScript = getKtFile(importedScript)?.declarations?.firstIsInstanceOrNull<KtScript>()
+                ?: return null
+            return resolveSession.getScriptDescriptor(ktScript) as ScriptDescriptor
+        }
 
-            fun errorDescriptor(errorDiagnostic: DiagnosticFactory1<PsiElement, String>?): ScriptDescriptor? {
+        private fun getKtFile(script: SourceCode): KtFile? {
-                reportErrorString1(errorDiagnostic, importedScriptFile.path)
+            if (script is KtFileScriptSource) return script.ktFile
+
+            // TODO: support any kind of ScriptSource.
+            if (script !is FileBasedScriptSource) return null
+
+            fun errorKtFile(errorDiagnostic: DiagnosticFactory1<PsiElement, String>?): KtFile? {
+                reportErrorString1(errorDiagnostic, script.file.path)
                 return null
             }
 
-            val vfile = localFS.findFileByPath(importedScriptFile.absolutePath)
+            val virtualFile = when (script) {
-                ?: return errorDescriptor(MISSING_IMPORTED_SCRIPT_FILE)
+                is VirtualFileScriptSource -> script.virtualFile
-            val psiFile = psiManager.findFile(vfile)
+                else -> localFS.findFileByPath(script.file.absolutePath) ?: return errorKtFile(MISSING_IMPORTED_SCRIPT_FILE)
-                ?: return errorDescriptor(MISSING_IMPORTED_SCRIPT_PSI)
+            }
-            // Note: is not an error now - if import references other valid source file, it is simply compiled along with script
+
-            // TODO: check if this is the behavior we want to have - see #KT-28916
+            val psiFile = psiManager.findFile(virtualFile) ?: return errorKtFile(MISSING_IMPORTED_SCRIPT_PSI)
-            val ktScript = (psiFile as? KtFile)?.declarations?.firstIsInstanceOrNull<KtScript>()
+            return psiFile as? KtFile
-                ?: return null
-            return resolveSession.getScriptDescriptor(ktScript) as ScriptDescriptor
         }
     }
 

plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/resolve/refineCompilationConfiguration.kt

@@ -114,7 +114,7 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
     abstract val dependenciesSources: List<File>
     abstract val javaHome: File?
     abstract val defaultImports: List<String>
-    abstract val importedScripts: List<File>
+    abstract val importedScripts: List<SourceCode>
 
     override fun equals(other: Any?): Boolean = script == (other as? ScriptCompilationConfigurationWrapper)?.script
 
@@ -141,9 +141,8 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
         override val defaultImports: List<String>
             get() = configuration?.get(ScriptCompilationConfiguration.defaultImports).orEmpty()
 
-        override val importedScripts: List<File>
+        override val importedScripts: List<SourceCode>
-            get() = configuration?.get(ScriptCompilationConfiguration.importScripts)
+            get() = configuration?.get(ScriptCompilationConfiguration.importScripts).orEmpty()
-                ?.mapNotNull { (it as? FileBasedScriptSource)?.file }.orEmpty()
 
         @Suppress("OverridingDeprecatedMember")
         override val legacyDependencies: ScriptDependencies?
@@ -178,8 +177,8 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
         override val defaultImports: List<String>
             get() = legacyDependencies?.imports.orEmpty()
 
-        override val importedScripts: List<File>
+        override val importedScripts: List<SourceCode>
-            get() = legacyDependencies?.scripts.orEmpty()
+            get() = legacyDependencies?.scripts?.map { FileScriptSource(it) }.orEmpty()
 
         override val configuration: ScriptCompilationConfiguration?
             get() {

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/dependencies/ScriptsCompilationDependencies.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import java.io.File
+import kotlin.script.experimental.host.FileBasedScriptSource
 
 data class ScriptsCompilationDependencies(
     val classpath: List<File>,
@@ -45,8 +46,10 @@ fun collectScriptsCompilationDependencies(
                 if (refinedConfiguration != null) {
                     collectedClassPath.addAll(refinedConfiguration.dependenciesClassPath)
 
-                    val sourceDependenciesRoots = refinedConfiguration.importedScripts.map {
+                    val sourceDependenciesRoots = refinedConfiguration.importedScripts.mapNotNull {
-                        KotlinSourceRoot(it.absolutePath, false)
+                        // TODO: support any kind of ScriptSource.
+                        val path = (it as? FileBasedScriptSource)?.file?.path ?: return@mapNotNull null
+                        KotlinSourceRoot(path, false)
                     }
                     val sourceDependencies =
                         createSourceFilesFromSourceRoots(