Scripting: avoid creating multiple PSIs for every imported script

#KT-53009 fixed #KT-42810 should be possible now too, but more testing is needed #KT-42101 can also be addressed now, but first the serialization of the imported scripts property should be solved
2022-07-20 19:42:53 +02:00
parent 9bc057afb2
commit ceea563d63
5 changed files with 90 additions and 42 deletions
@@ -5,20 +5,15 @@

 package org.jetbrains.kotlin.scripting.resolve

-import com.intellij.openapi.vfs.StandardFileSystems
-import com.intellij.openapi.vfs.VirtualFileManager
-import com.intellij.openapi.vfs.VirtualFileSystem
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.descriptors.annotations.FilteredAnnotations
 import org.jetbrains.kotlin.descriptors.impl.ClassConstructorDescriptorImpl
-import org.jetbrains.kotlin.descriptors.impl.PropertyDescriptorImpl
 import org.jetbrains.kotlin.descriptors.impl.ValueParameterDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.DiagnosticFactory1
 import org.jetbrains.kotlin.diagnostics.Errors
-import org.jetbrains.kotlin.diagnostics.Errors.MISSING_IMPORTED_SCRIPT_FILE
 import org.jetbrains.kotlin.diagnostics.Errors.MISSING_IMPORTED_SCRIPT_PSI
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -41,7 +36,10 @@ import org.jetbrains.kotlin.resolve.scopes.LexicalScopeImpl
 import org.jetbrains.kotlin.resolve.scopes.LexicalScopeKind
 import org.jetbrains.kotlin.resolve.scopes.utils.addImportingScope
 import org.jetbrains.kotlin.resolve.source.toSourceElement
-import org.jetbrains.kotlin.scripting.definitions.*
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
+import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.definitions.ScriptPriorities
+import org.jetbrains.kotlin.scripting.definitions.findScriptCompilationConfiguration
 import org.jetbrains.kotlin.types.TypeSubstitutor
 import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.types.typeUtil.isNothing
@@ -50,7 +48,6 @@ import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 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
@@ -185,10 +182,6 @@ class LazyScriptDescriptor(

    private inner class ImportedScriptDescriptorsFinder {

-        val localFS: VirtualFileSystem by lazy(LazyThreadSafetyMode.PUBLICATION) {
-            val fileManager = VirtualFileManager.getInstance()
-            fileManager.getFileSystem(StandardFileSystems.FILE_PROTOCOL)
-        }
        val psiManager by lazy(LazyThreadSafetyMode.PUBLICATION) { PsiManager.getInstance(scriptInfo.script.project) }

        operator fun invoke(importedScript: SourceCode): ScriptDescriptor? {
@@ -202,20 +195,13 @@ class LazyScriptDescriptor(
        private fun getKtFile(script: SourceCode): KtFile? {
            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)
+                reportErrorString1(errorDiagnostic, script.locationId ?: script.name ?: "unknown script")
                return null
            }

-            val virtualFile = when (script) {
-                is VirtualFileScriptSource -> script.virtualFile
-                else -> localFS.findFileByPath(script.file.absolutePath) ?: return errorKtFile(MISSING_IMPORTED_SCRIPT_FILE)
-            }
-
-            val psiFile = psiManager.findFile(virtualFile) ?: return errorKtFile(MISSING_IMPORTED_SCRIPT_PSI)
+            val psiFile = (script as? VirtualFileScriptSource)?.let { psiManager.findFile(it.virtualFile) }
+                ?: return errorKtFile(MISSING_IMPORTED_SCRIPT_PSI)
            return psiFile as? KtFile
        }
    }
@@ -9,9 +9,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.text.StringUtil
-import com.intellij.openapi.vfs.CharsetToolkit
-import com.intellij.openapi.vfs.LocalFileSystem
-import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.openapi.vfs.*
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiManager
@@ -24,6 +22,7 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.scripting.definitions.KotlinScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.runReadAction
+import org.jetbrains.kotlin.scripting.scriptFileName
 import org.jetbrains.kotlin.scripting.withCorrectExtension
 import java.io.File
 import java.net.URL
@@ -38,6 +37,7 @@ import kotlin.script.experimental.jvm.*
 import kotlin.script.experimental.jvm.compat.mapToDiagnostics
 import kotlin.script.experimental.jvm.impl.toClassPathOrEmpty
 import kotlin.script.experimental.jvm.impl.toDependencies
+import kotlin.script.experimental.util.PropertiesCollection

 internal fun VirtualFile.loadAnnotations(
    acceptedAnnotations: List<KClass<out Annotation>>,
@@ -146,7 +146,7 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {
            get() = configuration?.get(ScriptCompilationConfiguration.defaultImports).orEmpty()

        override val importedScripts: List<SourceCode>
-            get() = configuration?.get(ScriptCompilationConfiguration.importScripts).orEmpty()
+            get() = (configuration?.get(ScriptCompilationConfiguration.resolvedImportScripts) ?: configuration?.get(ScriptCompilationConfiguration.importScripts)).orEmpty()

        @Suppress("OverridingDeprecatedMember", "OVERRIDE_DEPRECATION")
        override val legacyDependencies: ScriptDependencies?
@@ -217,12 +217,15 @@ abstract class ScriptCompilationConfigurationWrapper(val script: SourceCode) {

 typealias ScriptCompilationConfigurationResult = ResultWithDiagnostics<ScriptCompilationConfigurationWrapper>

+val ScriptCompilationConfigurationKeys.resolvedImportScripts by PropertiesCollection.key<List<SourceCode>>(isTransient = true)
+
@Suppress("DEPRECATION")
 fun refineScriptCompilationConfiguration(
    script: SourceCode,
    definition: ScriptDefinition,
    project: Project,
-    providedConfiguration: ScriptCompilationConfiguration? = null // if null - take from definition
+    providedConfiguration: ScriptCompilationConfiguration? = null, // if null - take from definition
+    knownVirtualFileSources: MutableMap<String, VirtualFileScriptSource>? = null
 ): ScriptCompilationConfigurationResult {
    // TODO: add location information on refinement errors
    val ktFileSource = script.toKtFileSource(definition, project)
@@ -236,6 +239,8 @@ fun refineScriptCompilationConfiguration(
        return compilationConfiguration.refineOnAnnotations(script, collectedData)
            .onSuccess {
                it.refineBeforeCompiling(script, collectedData)
+            }.onSuccess {
+                it.resolveImportsToVirtualFiles(knownVirtualFileSources)
            }.onSuccess {
                ScriptCompilationConfigurationWrapper.FromCompilationConfiguration(
                    ktFileSource,
@@ -294,6 +299,50 @@ private fun additionalClasspath(definition: ScriptDefinition): List<File> {
        ?: definition.hostConfiguration[ScriptingHostConfiguration.configurationDependencies].toClassPathOrEmpty())
 }

+fun ScriptCompilationConfiguration.resolveImportsToVirtualFiles(
+    knownFileBasedSources: MutableMap<String, VirtualFileScriptSource>?
+)
+: ResultWithDiagnostics<ScriptCompilationConfiguration> {
+    // the resolving is needed while CoreVirtualFS does not cache the files, so attempt to find vf and then PSI by path leads
+    // to different PSI files, which breaks mappings needed by script descriptor
+    // resolving only to virtual file allows to simplify serialization and maybe a bit more future proof
+
+    val localFS: VirtualFileSystem by lazy(LazyThreadSafetyMode.NONE) {
+        val fileManager = VirtualFileManager.getInstance()
+        fileManager.getFileSystem(StandardFileSystems.FILE_PROTOCOL)
+    }
+
+    val resolvedImports = get(ScriptCompilationConfiguration.importScripts)?.map { sourceCode ->
+        when (sourceCode) {
+            is VirtualFileScriptSource -> sourceCode
+            is FileBasedScriptSource -> {
+                val path = sourceCode.file.normalize().absolutePath
+                knownFileBasedSources?.get(path) ?: run {
+                    val virtualFile = localFS.findFileByPath(path)
+                        ?: return@resolveImportsToVirtualFiles makeFailureResult("Imported source file not found: ${sourceCode.file}".asErrorDiagnostics())
+                    VirtualFileScriptSource(virtualFile).also {
+                        knownFileBasedSources?.set(path, it)
+                    }
+                }
+            }
+
+            else -> {
+                // TODO: support knownFileBasedSources here as well
+                val scriptFileName = sourceCode.scriptFileName(sourceCode, this)
+                val virtualFile = ScriptLightVirtualFile(
+                    scriptFileName,
+                    sourceCode.locationId,
+                    sourceCode.text
+                )
+                VirtualFileScriptSource(virtualFile)
+            }
+        }
+    }
+
+    val updatedConfiguration = if (resolvedImports.isNullOrEmpty()) this else this.with { resolvedImportScripts(resolvedImports) }
+    return updatedConfiguration.asSuccess()
+}
+
 internal fun makeScriptContents(
    file: VirtualFile,
    legacyDefinition: KotlinScriptDefinition,
@@ -10,10 +10,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
-import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
-import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
-import org.jetbrains.kotlin.scripting.resolve.ScriptReportSink
-import org.jetbrains.kotlin.scripting.resolve.refineScriptCompilationConfiguration
+import org.jetbrains.kotlin.scripting.resolve.*
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.read
 import kotlin.concurrent.write
@@ -22,6 +19,7 @@ import kotlin.script.experimental.api.ScriptCompilationConfiguration
 class CliScriptDependenciesProvider(project: Project) : ScriptDependenciesProvider(project) {
    private val cacheLock = ReentrantReadWriteLock()
    private val cache = hashMapOf<String, ScriptCompilationConfigurationResult?>()
+    private val knownVirtualFileSources = mutableMapOf<String, VirtualFileScriptSource>()

    override fun getScriptConfigurationResult(file: KtFile): ScriptCompilationConfigurationResult? = cacheLock.read {
        calculateRefinedConfiguration(file, null)
@@ -43,7 +41,10 @@ class CliScriptDependenciesProvider(project: Project) : ScriptDependenciesProvid
        else {
            val scriptDef = file.findScriptDefinition()
            if (scriptDef != null) {
-                val result = refineScriptCompilationConfiguration(KtFileScriptSource(file), scriptDef, project, providedConfiguration)
+                val result =
+                    refineScriptCompilationConfiguration(
+                        KtFileScriptSource(file), scriptDef, project, providedConfiguration, knownVirtualFileSources
+                    )

                ServiceManager.getService(project, ScriptReportSink::class.java)?.attachReports(file.virtualFile, result.reports)

@@ -6,12 +6,17 @@
 package org.jetbrains.kotlin.scripting.compiler.plugin.dependencies

 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiManager
 import org.jetbrains.kotlin.cli.common.config.KotlinSourceRoot
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.jvm.compiler.createSourceFilesFromSourceRoots
+import org.jetbrains.kotlin.cli.jvm.compiler.report
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
+import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
 import java.io.File
 import kotlin.script.experimental.api.ResultWithDiagnostics
 import kotlin.script.experimental.api.ScriptCompilationConfiguration
@@ -42,6 +47,7 @@ fun collectScriptsCompilationDependencies(
    var remainingSources = initialSources
    val knownSourcePaths = initialSources.mapNotNullTo(HashSet()) { it.virtualFile?.path }
    val importsProvider = ScriptDependenciesProvider.getInstance(project)
+    val psiManager by lazy(LazyThreadSafetyMode.NONE) { PsiManager.getInstance(project) }
    if (importsProvider != null) {
        while (true) {
            val newRemainingSources = ArrayList<KtFile>()
@@ -54,17 +60,23 @@ fun collectScriptsCompilationDependencies(
                    is ResultWithDiagnostics.Success -> {
                        collectedClassPath.addAll(refinedConfiguration.value.dependenciesClassPath)

-                        val sourceDependenciesRoots = refinedConfiguration.value.importedScripts.mapNotNull {
-                            // TODO: support any kind of ScriptSource.
-                            val path = (it as? FileBasedScriptSource)?.file?.absolutePath ?: return@mapNotNull null
-                            KotlinSourceRoot(path, false)
+                        val sourceDependencies = refinedConfiguration.value.importedScripts.mapNotNull {
+                            if (it is KtFileScriptSource) it.ktFile
+                            else {
+                                val virtualFileSource = (it as? VirtualFileScriptSource)
+                                    ?: error("expecting script sources resolved to virtual files here")
+                                (psiManager.findFile(virtualFileSource.virtualFile) as? KtFile).also {
+                                    if (it == null) {
+                                        configuration.report(
+                                            CompilerMessageSeverity.ERROR,
+                                            "imported file is not kotlin source: ${virtualFileSource.virtualFile.path}",
+                                            // TODO: consider receiving and using precise location from the resolver in the future
+                                            CompilerMessageLocation.create(source.virtualFile.path)
+                                        )
+                                    }
+                                }
+                            }
                        }
-                        val sourceDependencies =
-                            createSourceFilesFromSourceRoots(
-                                configuration, project, sourceDependenciesRoots,
-                                // TODO: consider receiving and using precise location from the resolver in the future
-                                source.virtualFile?.path?.let { CompilerMessageLocation.create(it) }
-                            )
                        if (sourceDependencies.isNotEmpty()) {
                            collectedSourceDependencies.add(
                                ScriptsCompilationDependencies.SourceDependencies(